from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch
from torch.utils.data import DataLoader

from .data import CurEDataset, validate_array_shape, validate_timestamp_shape
from .interpolation import interpolate_batch
from .metrics import masked_mse
from .modules import ITINCore


@dataclass
class CurEConfig:
    n_steps: int
    n_features: int
    rnn_hidden_size: int = 128
    batch_size: int = 32
    epochs: int = 30
    patience: int = 3
    learning_rate: float = 1e-3
    alpha: float = 1.2
    device: str | None = None


class CurEImputer:
    def __init__(
        self,
        n_steps: int,
        n_features: int,
        rnn_hidden_size: int = 128,
        batch_size: int = 32,
        epochs: int = 30,
        patience: int = 3,
        learning_rate: float = 1e-3,
        alpha: float = 1.2,
        device: str | None = None,
    ) -> None:
        self.config = CurEConfig(
            n_steps=n_steps,
            n_features=n_features,
            rnn_hidden_size=rnn_hidden_size,
            batch_size=batch_size,
            epochs=epochs,
            patience=patience,
            learning_rate=learning_rate,
            alpha=alpha,
            device=device,
        )
        self.device = torch.device(
            device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.model = ITINCore(n_steps, n_features, rnn_hidden_size).to(self.device)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
        self.best_state_dict: dict[str, torch.Tensor] | None = None
        self.best_loss = float("inf")

    def _assemble_training_input(self, batch: list) -> dict:
        (
            indices,
            X,
            missing_mask,
            timestamps,
            deltas,
            back_X,
            back_missing_mask,
            back_timestamps,
            back_deltas,
            *rest,
        ) = [item.to(self.device) if isinstance(item, torch.Tensor) else item for item in batch]

        inputs = {
            "indices": indices,
            "forward": {
                "X": X,
                "missing_mask": missing_mask,
                "timestamps": timestamps,
                "deltas": deltas,
            },
            "backward": {
                "X": back_X,
                "missing_mask": back_missing_mask,
                "timestamps": back_timestamps,
                "deltas": back_deltas,
            },
        }
        if len(rest) == 2:
            inputs["X_ori"] = rest[0]
            inputs["indicating_mask"] = rest[1]
        return inputs

    def _cur_e_regularization(self, inputs: dict, results: dict) -> torch.Tensor:
        X = inputs["forward"]["X"].detach().cpu().numpy().copy()
        missing_mask = inputs["forward"]["missing_mask"].detach().cpu().numpy()
        timestamps = inputs["forward"]["timestamps"].detach().cpu().numpy()
        X[missing_mask == 0] = np.nan
        interpolated, interpolation_valid_mask = interpolate_batch(
            X,
            timestamps=timestamps,
            return_valid_mask=True,
        )
        interpolated_t = torch.from_numpy(interpolated).to(self.device)
        interpolation_valid_mask_t = torch.from_numpy(interpolation_valid_mask).to(self.device)
        target_mask = (1.0 - inputs["forward"]["missing_mask"]) * interpolation_valid_mask_t
        return masked_mse(results["imputed_data"], interpolated_t, target_mask)

    def fit(
        self,
        train_X: np.ndarray,
        train_timestamps: np.ndarray | None = None,
        val_X: np.ndarray | None = None,
        val_timestamps: np.ndarray | None = None,
        val_X_ori: np.ndarray | None = None,
        val_indicating_mask: np.ndarray | None = None,
        verbose: bool = False,
    ) -> "CurEImputer":
        train_X = validate_array_shape(train_X)
        if train_timestamps is not None:
            train_timestamps = validate_timestamp_shape(train_timestamps, train_X)
        train_loader = DataLoader(
            CurEDataset(train_X, timestamps=train_timestamps),
            batch_size=self.config.batch_size,
            shuffle=True,
        )

        val_loader = None
        if val_X is not None:
            if val_X_ori is None or val_indicating_mask is None:
                raise ValueError("Validation requires val_X_ori and val_indicating_mask.")
            if val_timestamps is not None:
                val_timestamps = validate_timestamp_shape(val_timestamps, val_X)
            val_loader = DataLoader(
                CurEDataset(
                    val_X,
                    timestamps=val_timestamps,
                    X_ori=val_X_ori,
                    indicating_mask=val_indicating_mask,
                ),
                batch_size=self.config.batch_size,
                shuffle=False,
            )

        remaining_patience = self.config.patience
        for epoch in range(1, self.config.epochs + 1):
            self.model.train()
            train_losses = []
            for batch in train_loader:
                inputs = self._assemble_training_input(batch)
                self.optimizer.zero_grad()
                results = self.model(inputs, training=True)
                loss = results["loss"]
                loss = loss + self.config.alpha * self._cur_e_regularization(inputs, results)
                loss.backward()
                self.optimizer.step()
                train_losses.append(loss.item())

            monitor_loss = float(np.mean(train_losses))
            val_loss = None
            if val_loader is not None:
                self.model.eval()
                val_losses = []
                with torch.no_grad():
                    for batch in val_loader:
                        inputs = self._assemble_training_input(batch)
                        results = self.model(inputs, training=False)
                        val_losses.append(
                            masked_mse(
                                results["imputed_data"],
                                inputs["X_ori"],
                                inputs["indicating_mask"],
                            ).item()
                        )
                val_loss = float(np.mean(val_losses))
                monitor_loss = val_loss

            if verbose:
                if val_loss is None:
                    print(f"Epoch {epoch:03d} | train_loss={np.mean(train_losses):.6f}")
                else:
                    print(
                        f"Epoch {epoch:03d} | train_loss={np.mean(train_losses):.6f} | "
                        f"val_loss={val_loss:.6f}"
                    )

            if monitor_loss < self.best_loss:
                self.best_loss = monitor_loss
                self.best_state_dict = {
                    key: value.detach().cpu().clone() for key, value in self.model.state_dict().items()
                }
                remaining_patience = self.config.patience
            else:
                remaining_patience -= 1
                if remaining_patience <= 0:
                    break

        if self.best_state_dict is not None:
            self.model.load_state_dict(self.best_state_dict)
        self.model.eval()
        return self

    def predict(self, X: np.ndarray, timestamps: np.ndarray | None = None) -> np.ndarray:
        X = validate_array_shape(X)
        if timestamps is not None:
            timestamps = validate_timestamp_shape(timestamps, X)
        loader = DataLoader(
            CurEDataset(X, timestamps=timestamps),
            batch_size=self.config.batch_size,
            shuffle=False,
        )
        outputs = []
        self.model.eval()
        with torch.no_grad():
            for batch in loader:
                inputs = self._assemble_training_input(batch)
                results = self.model(inputs, training=False)
                outputs.append(results["imputed_data"].detach().cpu().numpy())
        return np.concatenate(outputs, axis=0)

    def fit_predict(
        self,
        train_X: np.ndarray,
        test_X: np.ndarray,
        train_timestamps: np.ndarray | None = None,
        test_timestamps: np.ndarray | None = None,
        val_X: np.ndarray | None = None,
        val_timestamps: np.ndarray | None = None,
        val_X_ori: np.ndarray | None = None,
        val_indicating_mask: np.ndarray | None = None,
        verbose: bool = False,
    ) -> np.ndarray:
        self.fit(
            train_X=train_X,
            train_timestamps=train_timestamps,
            val_X=val_X,
            val_timestamps=val_timestamps,
            val_X_ori=val_X_ori,
            val_indicating_mask=val_indicating_mask,
            verbose=verbose,
        )
        return self.predict(test_X, timestamps=test_timestamps)

    def save(self, path: str) -> None:
        payload = {
            "config": self.config.__dict__,
            "state_dict": self.model.state_dict(),
        }
        torch.save(payload, path)

    @classmethod
    def load(cls, path: str, map_location: str | None = None) -> "CurEImputer":
        payload = torch.load(path, map_location=map_location or "cpu")
        model = cls(**payload["config"])
        model.model.load_state_dict(payload["state_dict"])
        model.model.to(model.device)
        model.model.eval()
        model.best_state_dict = payload["state_dict"]
        return model
