import warnings
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from pydantic import BaseModel

warnings.filterwarnings("ignore", message="h5py not installed")

from typing import List

# Supervised Models
from optrade.torch.models.patchtst.model import PatchTST
from optrade.torch.models.recurrent.model import RecurrentModel
from optrade.torch.models.linear.model import Linear
from optrade.torch.models.dlinear.model import DLinear
from optrade.torch.models.tsmixer.model import TSMixer
from optrade.torch.models.emforecaster.model import EMForecaster
from optrade.torch.models.utils.revin import RevIN

# Optimizers and Schedulers
from torch import optim

# from time_series.utils.schedulers import WarmupCosineSchedule, PatchTSTSchedule
# from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR


def get_model(
    args: BaseModel,
    generator=torch.Generator(),
    target_channels_idx: List[int] = None,
) -> nn.Module:
    if args.exp.model_id == "PatchTST":
        model = PatchTST(
            num_enc_layers=args.patchtst.num_enc_layers,
            d_model=args.patchtst.d_model,
            d_ff=args.patchtst.d_ff,
            num_heads=args.patchtst.num_heads,
            num_channels=args.data.num_channels,
            seq_len=args.data.seq_len,
            pred_len=args.data.pred_len,
            attn_dropout=args.patchtst.attn_dropout,
            ff_dropout=args.patchtst.ff_dropout,
            pred_dropout=args.patchtst.pred_dropout,
            batch_first=args.train.batch_first,
            norm_mode=args.train.norm_mode,
            revin=args.train.revin,
            revout=args.train.revout,
            revin_affine=args.train.revin_affine,
            eps_revin=args.train.eps_revin,
            patch_dim=args.data.patch_dim,
            stride=args.data.patch_stride,
            return_head=args.train.return_head,
            head_type=args.train.head_type,
            channel_independent=args.train.channel_independent,  # Head only
            target_channels=target_channels_idx,  # Head only
        )
    elif args.exp.model_id == "RecurrentModel":
        model = RecurrentModel(
            d_model=args.rnn.d_model,
            num_enc_layers=args.rnn.num_enc_layers,
            pred_len=args.data.pred_len,
            backbone_id=args.rnn.backbone_id,
            bidirectional=args.rnn.bidirectional,
            dropout=args.train.dropout,
            seq_len=args.data.seq_len,
            patching=args.data.patching,
            patch_dim=args.data.patch_dim,
            patch_stride=args.data.patch_stride,
            num_channels=args.data.num_channels,
            head_type=args.train.head_type,
            norm_mode=args.train.norm_mode,
            revin=args.train.revin,
            revout=args.train.revout,
            revin_affine=args.train.revin_affine,
            eps_revin=args.train.eps_revin,
            last_state=args.rnn.last_state,
            avg_state=args.rnn.avg_state,
            return_head=args.train.return_head,
            channel_independent=args.train.channel_independent,
            target_channels=target_channels_idx,
        )
    elif args.exp.model_id == "Linear":
        model = Linear(
            seq_len=args.data.seq_len,
            pred_len=args.data.pred_len,
            num_channels=args.data.num_channels,
            norm_mode=args.train.norm_mode,
            revin=args.train.revin,
            revout=args.train.revout,
            revin_affine=args.train.revin_affine,
            eps_revin=args.train.eps_revin,
            channel_independent=args.train.channel_independent,
            target_channels=target_channels_idx,
        )
    elif args.exp.model_id == "DLinear":
        model = DLinear(
            task=args.exp.task,
            seq_len=args.data.seq_len,
            pred_len=args.data.pred_len,
            num_channels=args.data.num_channels,
            num_classes=args.data.num_classes,
            moving_avg=args.dlinear.moving_avg,
            individual=args.dlinear.individual,
            return_head=args.train.return_head,
            revin=args.train.revin,
            revout=args.train.revout,
            revin_affine=args.train.revin_affine,
            eps_revin=args.train.eps_revin,
            target_channels=target_channels_idx,
        )
    elif args.exp.model_id == "EMForecaster":
        model = EMForecaster(
            args=args,
            seed=args.exp.seed,
            seq_len=args.data.seq_len,
            pred_len=args.data.pred_len,
            num_channels=args.data.num_channels,
            revin=args.train.revin,
            revout=args.train.revout,
            revin_affine=args.train.revin_affine,
            eps_revin=args.train.eps_revin,
            patch_model_id=args.emf.patch_model_id,
            patch_norm=args.emf.patch_norm,
            patch_act=args.emf.patch_act,
            patch_dim=args.data.patch_dim,
            patch_stride=args.data.patch_stride,
            patch_embed_dim=args.emf.patch_embed_dim,
            pos_enc=args.emf.pos_enc,
            return_head=args.train.return_head,
            target_channels=target_channels_idx,
        )
    elif args.exp.model_id == "TSMixer":
        model = TSMixer(
            seq_len=args.data.seq_len,
            pred_len=args.data.pred_len,
            num_enc_layers=args.tsmixer.num_enc_layers,
            d_model=args.tsmixer.d_model,
            num_channels=args.data.num_channels,
            dropout=args.train.dropout,
            revin=args.train.revin,
            revin_affine=args.train.revin_affine,
            revout=args.train.revout,
            eps_revin=args.train.eps_revin,
            return_head=args.train.return_head,
            target_channels=target_channels_idx,
            channel_independent=args.train.channel_independent,
        )
    else:
        raise ValueError("Please select a valid model_id.")
    return model


def get_optim(args, model, optimizer_type="adamw"):
    # if args.exp.sklearn:
    #     return None

    optimizer_classes = {"adam": optim.Adam, "adamw": optim.AdamW}

    if optimizer_type not in optimizer_classes:
        raise ValueError("Please select a valid optimizer.")
    optimizer_class = optimizer_classes[optimizer_type]

    param_groups = exclude_weight_decay(
        model, args
    )  # Exclude bias and normalization parameters from weight decay

    optimizer = optimizer_class(param_groups)  # Set optimizer

    return optimizer


def exclude_weight_decay(model, args):
    # Separate parameters into those that will use weight decay and those that won't
    decay_params = []
    no_decay_params = []

    for name, param in model.named_parameters():
        if param.requires_grad:
            if "bias" in name or isinstance(
                param, (nn.BatchNorm1d, nn.BatchNorm2d, nn.LayerNorm)
            ):
                no_decay_params.append(param)
            else:
                decay_params.append(param)

    # Create parameter groups
    param_groups = [
        {"params": decay_params, "weight_decay": args.train.weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]

    return param_groups


def get_scheduler(args, scheduler_type, training_mode, optimizer, num_batches=0):
    # if args.exp.sklearn:
    #     return None

    if scheduler_type == "cosine":
        scheduler = CosineAnnealingLR(
            optimizer,
            T_max=args.train.epochs,
            eta_min=args.train.lr * 1e-2,
            last_epoch=args.scheduler.last_epoch,
        )
    elif scheduler_type == "patchtst" and training_mode == "supervised":
        scheduler = PatchTSTSchedule(optimizer, args, num_batches)
    elif scheduler_type == "onecycle" and training_mode == "supervised":
        scheduler = OneCycleLR(
            optimizer=optimizer,
            steps_per_epoch=num_batches,
            pct_start=args.scheduler.pct_start,
            epochs=args.train.epochs,
            max_lr=args.train.lr,
        )
    elif scheduler_type is None:
        return None
    else:
        raise ValueError("Please select a valid scheduler_type.")
    return scheduler


def get_criterion(args, criterion_type):
    if criterion_type == "MSE":
        return nn.MSELoss()
    elif criterion_type == "SmoothL1":
        return nn.SmoothL1Loss()
    elif criterion_type == "BCE":
        return nn.BCEWithLogitsLoss()
    elif criterion_type == "BCE_normal":
        return nn.BCELoss()
    elif criterion_type == "CE":
        return nn.CrossEntropyLoss()
    else:
        raise ValueError("Please select a valid criterion_type.")


def forward_pass(args, model, batch, model_id, device):
    if model_id in {
        "PatchTST",
        "TSMixer",
        "RecurrentModel",
        "Linear",
        "DLinear",
        "ModernTCN",
        "TimesNet",
        "EMForecaster",
    }:
        x = batch[0]
        x = x.to(device)

        if args.data.difference_input:
            x = x.diff(dim=1)
            # Pad the last dimension with the last value of the sequence
            x = torch.cat([x, x[:, -1].unsqueeze(1)], dim=1)

        output = model(x)
    else:
        raise ValueError("Please select a valid model_id.")

    return output


def compute_loss(output, batch, criterion, model_id, args, device):
    if model_id in {
        "PatchTST",
        "TSMixer",
        "RecurrentModel",
        "Linear",
        "DLinear",
        "ModernTCN",
        "TimesNet",
        "EMForecaster",
    }:
        target = batch[1].to(device)
        loss = criterion(output, target)
    else:
        raise ValueError("Please select a valid model_id.")

    return loss


def model_update(model, loss, optimizer, model_id):
    if model_id in {
        "PatchTST",
        "TSMixer",
        "RecurrentModel",
        "Linear",
        "DLinear",
        "ModernTCN",
        "TimesNet",
        "EMForecaster",
    }:
        loss.backward()
        optimizer.step()
    else:
        raise ValueError(f"{model_id} is not a valid model_id.")


def check_gradients(model, threshold_low=1e-5, threshold_high=1e2):
    vanishing = []
    exploding = []
    normal = []

    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norm = param.grad.norm().item()
            if grad_norm < threshold_low:
                vanishing.append((name, grad_norm))
            elif grad_norm > threshold_high:
                exploding.append((name, grad_norm))
            else:
                normal.append((name, grad_norm))

    print(f"Gradient statistics:")
    print(
        f"  Total parameters with gradients: {len(vanishing) + len(exploding) + len(normal)}"
    )
    print(
        f"  Vanishing gradients: {len(vanishing)} ({len(vanishing) / (len(vanishing) + len(exploding) + len(normal)) * 100:.2f}%)"
    )
    print(
        f"  Exploding gradients: {len(exploding)} ({len(exploding) / (len(vanishing) + len(exploding) + len(normal)) * 100:.2f}%)"
    )
    print(
        f"  Normal gradients: {len(normal)} ({len(normal) / (len(vanishing) + len(exploding) + len(normal)) * 100:.2f}%)"
    )

    if vanishing:
        print("\nVanishing gradients:")
        for name, grad_norm in vanishing[:10]:  # Print first 10
            print(f"  {name}: {grad_norm}")
        if len(vanishing) > 10:
            print(f"  ... and {len(vanishing) - 10} more")

    if exploding:
        print("\nExploding gradients:")
        for name, grad_norm in exploding[:10]:  # Print first 10
            print(f"  {name}: {grad_norm}")
        if len(exploding) > 10:
            print(f"  ... and {len(exploding) - 10} more")

    # Compute and print gradient statistics
    all_grads = [
        param.grad.norm().item()
        for name, param in model.named_parameters()
        if param.grad is not None
    ]
    if all_grads:
        print("\nGradient norm statistics:")
        print(f"  Mean: {np.mean(all_grads):.6f}")
        print(f"  Median: {np.median(all_grads):.6f}")
        print(f"  Std: {np.std(all_grads):.6f}")
        print(f"  Min: {np.min(all_grads):.6f}")
        print(f"  Max: {np.max(all_grads):.6f}")
