from typing import Any, Callable, Optional, Tuple
import torch
import torch.nn.functional as F
from torch import nn, optim
import lightning.pytorch as pl
import torchvision.models.video as tvmv
import sklearn.metrics as skm


"""Head model for SYNTAX prediction."""
class SyntaxLightningModule(pl.LightningModule):
    def __init__(
        self,
        num_classes,
        lr: float,
        variant: str,
        weight_decay: float = 0,
        max_epochs: int = None,
        weight_path: str = None,
        save_path: str = None,
        pl_weight_path: str = None,
        pt_weights_format: bool = False,
        sigma_a: float = 0,
        sigma_b: float = 1,
        **kwargs,
    ):
        self.save_hyperparameters()
        super().__init__()
        self.num_classes = num_classes
        self.save_path = save_path
        self.weight_path = weight_path
        self.variant = variant
        self.sigma_a = sigma_a
        self.sigma_b = sigma_b

        self.model = tvmv.r3d_18(weights=tvmv.R3D_18_Weights.DEFAULT)

        self.lr = lr
        self.loss_clf = nn.BCEWithLogitsLoss(reduction="none")
        self.loss_reg = nn.MSELoss(reduction="none")

        in_features = self.model.fc.in_features
        self.model.fc = nn.Linear(in_features=in_features, out_features=2, bias=True)

        if weight_path is not None:
            print("Load model weights (backbone)")
            self.load_weights_backbone(weight_path, self.model)

        if self.variant != "mean_out":
            self.model.fc = nn.Identity()

        if self.variant == "mean_out":
            pass
        elif self.variant in ("gru_mean", "gru_last"):
            self.rnn = nn.GRU(in_features, in_features // 4, batch_first=True)
            self.dropout = nn.Dropout(0.2)
            self.fc = nn.Linear(in_features=in_features // 4, out_features=num_classes, bias=True)
        elif self.variant in ("lstm_mean", "lstm_last"):
            self.lstm = nn.LSTM(
                input_size=in_features,
                hidden_size=in_features // 4,
                proj_size=num_classes,
                batch_first=True,
            )
        elif self.variant == "mean":
            self.fc = nn.Linear(in_features=in_features, out_features=num_classes, bias=True)
        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=in_features,
                nhead=4,
                batch_first=True,
                dim_feedforward=in_features // 4,
            )
            self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
            self.dropout = nn.Dropout(0.2)
            self.fc = nn.Linear(in_features=in_features, out_features=num_classes, bias=True)
            if self.variant == "bert_cls2":
                self.cls = nn.Parameter(torch.randn(1, 1, in_features))
        else:
            raise ValueError(f"Unknown model variant {self.variant}")

        if pl_weight_path is not None:
            print(f"Load LightningModule weights from {pl_weight_path}")

            if pt_weights_format:
                pl_state_dict = torch.load(pl_weight_path, weights_only=False)
            else:
                pl_state_dict = torch.load(pl_weight_path, weights_only=False)["state_dict"]

            self.load_weights(pl_state_dict, self.model, "model")

            if self.variant == "mean_out":
                pass
            elif self.variant in ("gru_mean", "gru_last"):
                self.load_weights(pl_state_dict, self.rnn, "rnn")
                self.load_weights(pl_state_dict, self.fc, "fc")
            elif self.variant in ("lstm_mean", "lstm_last"):
                self.load_weights(pl_state_dict, self.lstm, "lstm")
            elif self.variant == "mean":
                self.load_weights(pl_state_dict, self.fc, "fc")
            elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
                self.load_weights(pl_state_dict, self.encoder, "encoder")
                self.load_weights(pl_state_dict, self.fc, "fc")
                if self.variant == "bert_cls2":
                    old_shape = self.cls.shape
                    self.cls = nn.Parameter(pl_state_dict["cls"])
                    assert old_shape == self.cls.shape
            else:
                raise ValueError(f"Unknown model variant {self.variant}")

        self.max_epochs = max_epochs
        self.weight_decay = weight_decay

        self.y_val = []
        self.p_val = []
        self.r_val = []
        self.ty_val = []
        self.tp_val = []

    def load_weights_backbone(self, weight_path: str, model: nn.Module) -> None:
        """
                Universal loader for backbone weights (r3d_18).

                - If the file is a Lightning checkpoint (a dict with 'state_dict'),
                    extract state_dict['state_dict'] and strip the 'model.' prefix.
                - If the file is a raw state_dict (.pt/.pth) saved via model.state_dict(),
                    load it directly.

                Before loading, drop any keys whose tensor shapes do not match the module.
        """
        obj = torch.load(weight_path, weights_only=False, map_location="cpu")

        if isinstance(obj, dict) and "state_dict" in obj:
            raw_state = obj["state_dict"]
            state_dict = {k.replace("model.", ""): v for k, v in raw_state.items()}
            src_type = "lightning_checkpoint"
        else:
            state_dict = obj
            src_type = "raw_state_dict"

        current_state = model.state_dict()
        filtered_state = {}
        mismatched_keys = []

        for k, v in state_dict.items():
            if k in current_state and current_state[k].shape == v.shape:
                filtered_state[k] = v
            else:
                mismatched_keys.append(k)

        incompatible = model.load_state_dict(filtered_state, strict=False)

        loaded_keys = [k for k in filtered_state.keys() if k not in incompatible.missing_keys]
        print(
            f"[Backbone] Loaded weights from '{weight_path}' "
            f"(type={src_type}): {len(loaded_keys)} params, "
            f"missing={len(incompatible.missing_keys)}, "
            f"unexpected={len(incompatible.unexpected_keys)}, "
            f"skipped_mismatched={len(mismatched_keys)}"
        )
        if mismatched_keys:
            print(f"[Backbone] Size‑mismatched keys (skipped), example: {mismatched_keys[:5]}")
        if incompatible.missing_keys:
            print(f"[Backbone] Missing keys after filtering, example: {incompatible.missing_keys[:5]}")
        if incompatible.unexpected_keys:
            print(f"[Backbone] Unexpected keys after filtering, example: {incompatible.unexpected_keys[:5]}")

    def load_weights(self, state_dict, module, prefix: str):
        """Filter and load only the weights that belong to the given module."""
        module_state = {
            k.replace(f"{prefix}.", ""): v
            for k, v in state_dict.items()
            if k.startswith(prefix)
        }
        missing, unexpected = module.load_state_dict(module_state, strict=False)
        if missing:
            print(f"Missing keys for {prefix}: {missing}")
        if unexpected:
            print(f"Unexpected keys for {prefix}: {unexpected}")

    def forward(self, x):
        batch_seq_shape = x.shape[0:2]
        x = torch.flatten(x, start_dim=0, end_dim=1)
        x = self.model(x)
        x = torch.unflatten(x, 0, batch_seq_shape)

        if self.variant == "mean_out":
            x = torch.mean(x, dim=1)
        elif self.variant in ("gru_mean", "gru_last"):
            _all_outs_, [_last_out_] = self.rnn(x)
            if self.variant == "gru_mean":
                x = torch.mean(_all_outs_, dim=1)
            else:
                x = _last_out_
            x = self.dropout(x)
            x = self.fc(x)
        elif self.variant in ("lstm_mean", "lstm_last"):
            _all_outs_, (_last_out_, _last_state_) = self.lstm(x)
            if self.variant == "lstm_mean":
                x = torch.mean(_all_outs_, dim=1)
            else:
                x = _last_out_
        elif self.variant == "mean":
            x = torch.mean(x, dim=1)
            x = self.fc(x)
        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
            if self.variant == "bert_cls":
                x = F.pad(x, (0, 0, 1, 0), "constant", 0)
            elif self.variant == "bert_cls2":
                bs = x.size(0)
                x = torch.cat([self.cls.expand(bs, -1, -1), x], dim=1)
            x = self.encoder(x)
            if self.variant == "bert_mean":
                x = torch.mean(x, dim=1)
            else:
                x = x[:, 0, :]
            x = self.dropout(x)
            x = self.fc(x)
        else:
            raise ValueError(f"Unknown model variant {self.variant}")

        return x

    def training_step(self, batch, batch_idx):
        x, y, target, path = batch
        y_hat = self(x)
        yp_clf = y_hat[:, 0:1]
        yp_reg = y_hat[:, 1:]

        weights_clf = torch.where(y > 0, 1.0, 0.2)
        clf_loss = self.loss_clf(yp_clf, y)
        clf_loss = (clf_loss * weights_clf).mean()

        reg_loss_raw = self.loss_reg(yp_reg, target)
        sigma = self.sigma_a * target + self.sigma_b
        reg_loss = (reg_loss_raw / (sigma ** 2)).mean()

        loss = clf_loss + 0.5 * reg_loss

        y_pred = torch.sigmoid(yp_clf)
        y_bin = torch.round(y.cpu().detach()).int()
        y_pred_bin = torch.round(y_pred.cpu().detach()).int()

        self.log("train_clf_loss", clf_loss, prog_bar=True, sync_dist=True)
        self.log("train_val_loss", reg_loss, prog_bar=True, sync_dist=True)
        self.log("train_full_loss", loss, prog_bar=True, sync_dist=True)
        self.log("train_f1", skm.f1_score(y_bin, y_pred_bin, zero_division=0),
                 prog_bar=True, sync_dist=True)
        self.log("train_acc", skm.accuracy_score(y_bin, y_pred_bin),
                 prog_bar=True, sync_dist=True)

        return loss

    def validation_step(self, batch, batch_idx):
        x, y, target, path = batch
        y_hat = self(x)
        yp_clf = y_hat[:, 0:1]
        yp_reg = y_hat[:, 1:]

        loss = self.loss_clf(yp_clf, y)
        reg_loss_raw = self.loss_reg(yp_reg, target)
        loss = loss.mean()

        y_pred = torch.sigmoid(yp_clf)

        self.y_val.append(int(y[..., 0].cpu()))
        self.p_val.append(float(y_pred[..., 0].cpu()))
        self.r_val.append(round(float(y_pred[..., 0].cpu())))

        self.ty_val.append(float(target[..., 0].cpu()))
        self.tp_val.append(float(yp_reg[..., 0].cpu()))

        clf_loss = self.loss_clf(yp_clf, y)
        reg_loss_raw = self.loss_reg(yp_reg, target)
        sigma = self.sigma_a * target + self.sigma_b
        reg_loss = (reg_loss_raw / (sigma ** 2)).mean()

        loss = clf_loss + 0.5 * reg_loss

        return loss

    def on_validation_epoch_end(self):
        try:
            auc = skm.roc_auc_score(self.y_val, self.p_val)
            f1 = skm.f1_score(self.y_val, self.r_val, zero_division=0)
            acc = skm.accuracy_score(self.y_val, self.r_val)
            mae = skm.mean_absolute_error(self.y_val, self.r_val)
            self.log("val_auc", auc, prog_bar=True, sync_dist=True)
            self.log("val_f1", f1, prog_bar=True, sync_dist=True)
            self.log("val_acc", acc, prog_bar=True, sync_dist=True)
            self.log("val_mae", mae, prog_bar=True, sync_dist=True)

            rmse = skm.root_mean_squared_error(self.ty_val, self.tp_val)
            self.log("val_rmse", rmse, prog_bar=True, sync_dist=True)

        except ValueError as err:
            print(err)
            print("Y_VAL", self.y_val)
            print("P_VAL", self.p_val)
        self.y_val.clear()
        self.p_val.clear()
        self.r_val.clear()
        self.ty_val.clear()
        self.tp_val.clear()

    def on_train_epoch_end(self) -> None:
        self.log(
            "lr",
            self.optimizers().optimizer.param_groups[0]["lr"],
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

    def configure_optimizers(self):
        if self.weight_path:
            if self.variant == "mean_out":
                trainable_modules = [self.model.fc]
            elif self.variant in ("gru_mean", "gru_last"):
                trainable_modules = [self.rnn, self.fc]
            elif self.variant in ("lstm_mean", "lstm_last"):
                trainable_modules = [self.lstm]
            elif self.variant == "mean":
                trainable_modules = [self.fc]
            elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
                trainable_modules = [self.encoder, self.fc]
                if self.variant == "bert_cls2":
                    trainable_modules.append(self.cls)
            else:
                trainable_modules = []

            for param in self.parameters():
                param.requires_grad = False

            for m in trainable_modules:
                for p in m.parameters():
                    p.requires_grad = True

            params = [p for m in trainable_modules for p in m.parameters()]
        else:
            for param in self.parameters():
                param.requires_grad = True
            params = self.parameters()

        optimizer = optim.Adam(params, lr=self.lr, weight_decay=self.weight_decay)

        if self.max_epochs is not None:
            lr_scheduler = optim.lr_scheduler.OneCycleLR(
                optimizer=optimizer, max_lr=self.lr, total_steps=self.max_epochs
            )
            return [optimizer], [lr_scheduler]
        else:
            return optimizer

    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
        """Inference step."""
        x, y, target, path = batch
        y_hat = self(x)
        yp_clf = y_hat[:, 0:1]
        yp_reg = y_hat[:, 1:]
        y_pred = torch.sigmoid(yp_clf)

        return {
            "y": y,
            "y_pred": torch.round(y_pred),
            "y_prob": y_pred,
            "y_reg": yp_reg,
            "target": target,
        }