import os
import json
import joblib
import numpy as np
import optuna
import time
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.model_selection import train_test_split, StratifiedShuffleSplit
from sklearn.metrics import roc_auc_score
from sklearn.linear_model import LogisticRegression

import onnx
from onnx import helper, numpy_helper, TensorProto

from ...utils.logging import logger

from ... import ONNX_TARGET_OPSET, ONNX_IR_VERSION


MAX_NUM_TRIALS = 50

NUM_EPOCHS = 30
BATCH_SIZE = 32


class HeadNN(nn.Module):
    """
    HeadNN is a small neural network designed for binary classification tasks.
    It supports configurations with up to two hidden layers.
    """

    def __init__(self, input_dim, n_hidden, scale1, scale2, dropout):
        super().__init__()
        layers = []
        if n_hidden == 0:
            layers.append(nn.Linear(input_dim, 1))
        elif n_hidden == 1:
            h1 = max(1, int(input_dim * scale1))
            layers.extend(
                [
                    nn.Linear(input_dim, h1),
                    nn.ReLU(),
                    nn.Dropout(dropout),
                    nn.Linear(h1, 1),
                ]
            )
        else:
            h1 = max(1, int(input_dim * scale1))
            h2 = max(1, int(h1 * scale2))
            layers.extend(
                [
                    nn.Linear(input_dim, h1),
                    nn.ReLU(),
                    nn.Dropout(dropout),
                    nn.Linear(h1, h2),
                    nn.ReLU(),
                    nn.Dropout(dropout),
                    nn.Linear(h2, 1),
                ]
            )
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x).squeeze(-1)


def find_params(X, y, num_trials):
    """
    Run Optuna hyperparameter optimization for HeadNN.
    Evaluation metric: ROC AUC.
    """
    n_trials = min(num_trials, MAX_NUM_TRIALS)
    epochs = NUM_EPOCHS
    batch_size = BATCH_SIZE

    def objective(trial):
        input_dim = X.shape[1]
        n_hidden = trial.suggest_int("n_hidden", 0, 2)
        scale1 = trial.suggest_float("scale1", 0.1, 0.5)
        scale2 = trial.suggest_float("scale2", 0.1, 0.5)
        dropout = trial.suggest_float("dropout", 0.0, 0.5)
        lr = trial.suggest_float("lr", 1e-4, 1e-1, log=True)

        X_train, X_val, y_train, y_val = train_test_split(
            X, y, test_size=0.2, stratify=y
        )

        model = HeadNN(input_dim, n_hidden, scale1, scale2, dropout)
        optimizer = optim.Adam(model.parameters(), lr=lr)

        pos_weight = torch.tensor(
            [(len(y_train) - y_train.sum()) / y_train.sum()], dtype=torch.float32
        )
        loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

        X_train_t = torch.tensor(X_train, dtype=torch.float32)
        y_train_t = torch.tensor(y_train, dtype=torch.float32)
        X_val_t = torch.tensor(X_val, dtype=torch.float32)

        for epoch in range(epochs):
            model.train()
            for i in range(0, len(X_train_t), batch_size):
                xb = X_train_t[i : i + batch_size]
                yb = y_train_t[i : i + batch_size]
                optimizer.zero_grad()
                logits = model(xb)
                loss = loss_fn(logits, yb)
                loss.backward()
                optimizer.step()

        model.eval()
        with torch.no_grad():
            logits_val = model(X_val_t).cpu().numpy()
            preds_val = 1 / (1 + np.exp(-logits_val))
        auc = roc_auc_score(y_val, preds_val)
        return auc

    study = optuna.create_study(direction="maximize")
    study.enqueue_trial(
        {"n_hidden": 1, "scale1": 0.5, "scale2": 0.5, "dropout": 0.2, "lr": 1e-3}
    )
    study.optimize(objective, n_trials=n_trials)

    results = {
        "n_hidden": study.best_params["n_hidden"],
        "scale1": study.best_params["scale1"],
        "scale2": study.best_params["scale2"],
        "dropout": study.best_params["dropout"],
        "lr": study.best_params["lr"],
        "epochs": epochs,
        "batch_size": batch_size,
        "input_dim": X.shape[1],
    }

    return results


class Head(BaseEstimator, ClassifierMixin):
    """
    Binary classification head wrapping HeadNN, trained with BCEWithLogitsLoss and class weighting.
    """

    def __init__(
        self,
        input_dim,
        n_hidden=1,
        scale1=0.5,
        scale2=0.5,
        dropout=0.0,
        lr=1e-3,
        epochs=30,
        batch_size=32,
        device=None,
    ):
        self.input_dim = input_dim
        self.n_hidden = n_hidden
        self.scale1 = scale1
        self.scale2 = scale2
        self.dropout = dropout
        self.lr = lr
        self.epochs = epochs
        self.batch_size = batch_size
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None

    def _fit(self, X, y):
        logger.info("Fitting the MLP model...")
        self.model = HeadNN(
            self.input_dim, self.n_hidden, self.scale1, self.scale2, self.dropout
        ).to(self.device)

        pos_weight = torch.tensor(
            [(len(y) - y.sum()) / y.sum()], dtype=torch.float32
        ).to(self.device)

        loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
        optimizer = optim.Adam(self.model.parameters(), lr=self.lr)

        X_t = torch.tensor(X, dtype=torch.float32).to(self.device)
        y_t = torch.tensor(y, dtype=torch.float32).to(self.device)

        for _ in range(self.epochs):
            self.model.train()
            for i in range(0, len(X_t), self.batch_size):
                xb = X_t[i : i + self.batch_size]
                yb = y_t[i : i + self.batch_size]
                optimizer.zero_grad()
                logits = self.model(xb)
                loss = loss_fn(logits, yb)
                loss.backward()
                optimizer.step()
        return self

    def _calibrate(self, X, y):
        logger.info("Calibrating the model using Platt scaling...")
        splitter = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=42)
        y_hat = []
        y_true = []
        t0 = time.time()
        for train_idx, val_idx in splitter.split(X, y):
            X_train, X_val = X[train_idx], X[val_idx]
            y_train, y_val = y[train_idx], y[val_idx]
            self._fit(X_train, y_train)
            logits_val = self.predict_raw(X_val)
            y_hat += list(logits_val)
            y_true += list(y_val)
            t1 = time.time()
            if (t1 - t0) > 60:
                break
        y_hat = np.array(y_hat)
        logger.debug("Shape of y_hat: {}".format(y_hat.shape))
        y_true = np.array(y_true)
        logger.debug("Shape of y_true: {}".format(y_true.shape))
        self.calibrator = LogisticRegression(class_weight="balanced")
        self.calibrator.fit(y_hat.reshape(-1, 1), y_true)
        self.score = roc_auc_score(
            y_true, self.calibrator.predict_proba(y_hat.reshape(-1, 1))[:, 1]
        )
        logger.debug("Done with calibration! Score: {:.4f}".format(self.score))
        return self.score

    def fit(self, X, y):
        self._calibrate(X, y)
        self._fit(X, y)
        return self

    def predict_raw(self, X):
        self.model.eval()
        X_t = torch.tensor(X, dtype=torch.float32).to(self.device)
        with torch.no_grad():
            logits = self.model(X_t).cpu().numpy()
        return logits

    def predict_proba(self, X):
        logits = self.predict_raw(X)
        probs = self.calibrator.predict_proba(logits.reshape(-1, 1))
        return probs

    def predict(self, X):
        return (self.predict_proba(X)[:, 1] >= 0.5).astype(int)

    def save(self, name: str, model_dir: str):
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
        model_path = os.path.join(model_dir, f"{name}.pth")
        torch.save(self.model.state_dict(), model_path)
        metadata = {
            "input_dim": self.input_dim,
            "n_hidden": self.n_hidden,
            "scale1": self.scale1,
            "scale2": self.scale2,
            "dropout": self.dropout,
            "lr": self.lr,
            "epochs": self.epochs,
            "batch_size": self.batch_size,
            "device": self.device,
            "score": self.score,
        }
        meta_path = os.path.join(model_dir, f"{name}_metadata.json")
        with open(meta_path, "w") as f:
            json.dump(metadata, f)
        joblib.dump(
            self.calibrator, os.path.join(model_dir, f"{name}_calibrator.joblib")
        )

    @classmethod
    def load(cls, name: str, model_dir: str):
        meta_path = os.path.join(model_dir, f"{name}_metadata.json")
        with open(meta_path, "r") as f:
            metadata = json.load(f)
        input_dim = metadata["input_dim"]
        n_hidden = metadata["n_hidden"]
        scale1 = metadata["scale1"]
        scale2 = metadata["scale2"]
        dropout = metadata["dropout"]
        lr = metadata["lr"]
        epochs = metadata["epochs"]
        batch_size = metadata["batch_size"]
        device = metadata["device"]

        model_path = os.path.join(model_dir, f"{name}.pth")
        model = HeadNN(input_dim, n_hidden, scale1, scale2, dropout).to(device)
        model.load_state_dict(torch.load(model_path, map_location=device))

        obj = cls(
            input_dim, n_hidden, scale1, scale2, dropout, lr, epochs, batch_size, device
        )
        obj.model = model

        obj.calibrator = joblib.load(
            os.path.join(model_dir, f"{name}_calibrator.joblib")
        )
        obj.score = metadata.get("score", None)

        return obj


def convert_to_onnx(name: str, model_dir: str):
    """
    Export Torch head -> logits and append Platt calibrator in ONNX.
    Final output: flat 1D vector [batch_size].
    """
    head = Head.load(name, model_dir)
    model = head.model
    model = model.to("cpu")
    model.eval()

    onnx_path = os.path.join(model_dir, f"{name}.onnx")
    dummy_input = torch.randn(1, head.input_dim, dtype=torch.float32, device="cpu")

    # 1) Export Torch model (produces logits_{name})
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        input_names=[f"input_{name}"],
        output_names=[f"logits_{name}"],
        dynamic_axes={
            f"input_{name}": {0: "batch_size"},
            f"logits_{name}": {0: "batch_size"},
        },
        opset_version=ONNX_TARGET_OPSET,
    )

    # 2) Patch ONNX: add Platt = Mul + Add + Sigmoid (+ Reshape to 1D)
    onnx_model = onnx.load(onnx_path)

    coef = np.asarray(head.calibrator.coef_, dtype=np.float32).reshape(
        1,
    )  # scalar
    intercept = np.asarray(head.calibrator.intercept_, dtype=np.float32).reshape(
        1,
    )
    shape1d = np.array([-1], dtype=np.int64)

    onnx_model.graph.initializer.extend(
        [
            numpy_helper.from_array(coef, name=f"calib_coef_{name}"),
            numpy_helper.from_array(intercept, name=f"calib_intercept_{name}"),
            numpy_helper.from_array(shape1d, name=f"shape1d_{name}"),
        ]
    )

    mul_node = helper.make_node(
        "Mul",
        inputs=[f"logits_{name}", f"calib_coef_{name}"],
        outputs=[f"logits_scaled_{name}"],
        name=f"Calib_Mul_{name}",
    )
    add_node = helper.make_node(
        "Add",
        inputs=[f"logits_scaled_{name}", f"calib_intercept_{name}"],
        outputs=[f"logits_shifted_{name}"],
        name=f"Calib_Add_{name}",
    )
    sigmoid_node = helper.make_node(
        "Sigmoid",
        inputs=[f"logits_shifted_{name}"],
        outputs=[f"probs_2d_{name}"],
        name=f"Calib_Sigmoid_{name}",
    )
    reshape_node = helper.make_node(
        "Reshape",
        inputs=[f"probs_2d_{name}", f"shape1d_{name}"],
        outputs=[f"output_{name}"],
        name=f"Output_Reshape1D_{name}",
    )

    onnx_model.graph.node.extend([mul_node, add_node, sigmoid_node, reshape_node])

    # Replace graph outputs with our final 1D tensor
    del onnx_model.graph.output[:]
    onnx_model.graph.output.extend(
        [
            helper.make_tensor_value_info(
                f"output_{name}", TensorProto.FLOAT, ["batch_size"]
            )
        ]
    )

    onnx_model.graph.name = f"{name}"
    onnx_model.ir_version = ONNX_IR_VERSION

    onnx.checker.check_model(onnx_model)
    onnx.save(onnx_model, onnx_path)
    logger.info(f"ONNX model with calibrator saved to {onnx_path}")
    return onnx_path
