import json
from pathlib import Path

import numpy as np
import torch
from torch import nn


def load_config(path):
    with open(path, encoding="utf-8") as handle:
        return json.load(handle)


def write_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(value, handle, indent=2, sort_keys=True)


class StreamflowLSTM(nn.Module):
    """Paper stack with explicit inter-layer activation semantics."""

    def __init__(self, input_size=23, hidden_size=50, dropout=0.1):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.LSTM(input_size, hidden_size, batch_first=True),
            nn.LSTM(hidden_size, hidden_size, batch_first=True),
            nn.LSTM(hidden_size, hidden_size, batch_first=True),
        ])
        self.dropout = nn.Dropout(dropout)
        self.dense = nn.Linear(hidden_size, 1)

    def forward(self, inputs):
        if inputs.ndim != 3 or inputs.shape[-2:] != (28, 23):
            raise ValueError(f"expected [B,28,23], got {tuple(inputs.shape)}")
        values, _ = self.layers[0](inputs)
        values = self.dropout(torch.relu(values))
        values, _ = self.layers[1](values)
        values = self.dropout(torch.relu(values))
        values, _ = self.layers[2](values)
        values = torch.tanh(values[:, -1])
        return self.dense(values).squeeze(-1)


def nse(prediction, target):
    prediction, target = np.asarray(prediction), np.asarray(target)
    denominator = np.sum((target - target.mean()) ** 2)
    return float(1.0 - np.sum((prediction - target) ** 2) / max(denominator, 1e-12))


def metrics(prediction, target):
    prediction = np.asarray(prediction, dtype=np.float64)
    target = np.asarray(target, dtype=np.float64)
    correlation = float(np.corrcoef(prediction, target)[0, 1]) if prediction.size > 1 else 0.0
    if not np.isfinite(correlation):
        correlation = 0.0
    alpha = float(prediction.std() / max(target.std(), 1e-12))
    beta = float(prediction.mean() / max(target.mean(), 1e-12))
    kge = float(1.0 - np.sqrt((correlation - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2))
    return {
        "kge": kge,
        "kge_r": correlation,
        "kge_alpha": alpha,
        "kge_beta": beta,
        "nse": nse(prediction, target),
        "rmse_m3_s": float(np.sqrt(np.mean((prediction - target) ** 2))),
    }


def load_member(payload, device):
    model = StreamflowLSTM(hidden_size=payload["hidden_size"], dropout=payload["dropout"]).to(device)
    model.load_state_dict(payload["state_dict"])
    model.eval()
    return model, payload