"""
Configuration utilities for setting up device, data loaders, model, loss function,
optimizer, scheduler, and plotting based on a configuration dictionary.
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from functools import partial

from models.MANO import MANO
from models.Swin_MANO import SwinV2_MANO

from plot_utils.plot_NS import plot_NS_2D_trainer

from data.dataset_NS import Dataset2D
from data.dataloader_NS import DataLoader2D
from data.dataloader_IC import img_class_loader


def get_device(cfg):
    if cfg["device"] != "cpu":
        if torch.cuda.is_available():
            return cfg["device"]
        else:
            raise ValueError(
                f"{cfg['device']} is not available, cosider to change cfg['device'] to 'cpu' or connect to a cluster :)"
            )
    else:
        return "cpu"


def get_dataloader(cfg):
    if cfg["data_name"].lower() in {"darcy", "navier_stokes", "allen_cahn"}:
        dataset = Dataset2D(data_name=cfg["data_name"], res=cfg["resolution"])

        n_train_dict = {
            **{
                f"darcy_{r}": v
                for r, v in [
                    (16, 1000),
                    (32, 5000),
                    (64, 5000),
                    (128, 5000),
                    (412, 4000),
                ]
            }
        }

        if cfg["n_train"].lower() == "default":
            n_train = n_train_dict[f"{cfg['data_name'].lower()}_{cfg['resolution']}"]
        else:
            n_train = cfg["n_train"]

        data_loader = DataLoader2D(
            dataset,
            n_train,
            batch_size=cfg["batch_size_per_gpu"],
            stack_regular_grid=cfg["stack_regular_grid"],
        )
        train_loader, test_loader = (
            data_loader.train_dataloader(),
            data_loader.val_dataloader(),
        )
        return train_loader, test_loader

    elif cfg["data_name"].lower() in [
        "mnist",
        "cifar10",
        "cifar100",
        "food101",
        "flowers102",
    ]:
        train_loader, test_loader, input_shape, output_shape = img_class_loader(
            data_name=cfg["data_name"],
            batch_size=cfg["batch_size_per_gpu"],
            to_print=False,
            colored=False,
            visually=False,
            resolution=cfg["resolution"],
        )
        return train_loader, test_loader, input_shape, output_shape


def get_resolution(cfg):
    if cfg["resolution"] != "default":
        return cfg["resolution"]
    else:
        res_dict = {"mnist": 28, "cifar10": 32, "cifar100": 32}
        return res_dict[cfg["data_name"].lower()]


def get_model(cfg, device, output_shape=None, resolution=None):
    if cfg["model"].lower() == "mano":
        model = MANO(
            device=device,
            image_size=resolution,
            dim=cfg["dim"],
            depth=cfg["depth"],
            heads=cfg["heads"],
            dim_head=cfg["dim_head"],
            att_dropout=cfg["att_dropout"],
            channel_scale=cfg["channel_scale"],
            mlp_dim=cfg["mlp_dim"],
            channels=cfg["channels"],
            emb_dropout=cfg["emb_dropout"],
            local_attention_span=cfg["local_attention_span"],
            local_attention_stride=cfg["local_attention_stride"],
            att_sampling=cfg["att_sampling"],
            att_sampling_rate=cfg["att_sampling_rate"],
        )
    elif cfg["model"].lower() == "swin_mano":
        model = SwinV2_MANO(
            img_size=cfg["resolution"],
            patch_size=cfg["patch_size"],
            in_chans=cfg["in_channels"],
            num_classes=output_shape,
            global_pool="avg",
            embed_dim=cfg["embed_dim"],
            depths=cfg["depths"],
            num_heads=cfg["num_heads"],
            window_size=cfg["window_size"],
            always_partition=False,
            strict_img_size=True,
            mlp_ratio=cfg["mlp_ratio"],
            qkv_bias=cfg["qkv_bias"],
            drop_rate=cfg["drop_rate"],
            proj_drop_rate=cfg["proj_drop_rate"],
            attn_drop_rate=cfg["attn_drop_rate"],
            drop_path_rate=cfg["drop_path_rate"],
            act_layer=cfg["act_layer"],
            norm_layer=nn.LayerNorm,
            pretrained_window_sizes=cfg["pretrained_window_sizes"],
            sampling_rate=cfg["att_sampling_rate"],
            down_up_sampling=cfg["att_down_up_sampling"],
        )
    return model.to(device, non_blocking=True)


def get_loss_fn(cfg):
    if cfg["loss_fn"].lower() == "mse":
        return torch.nn.MSELoss()


def get_optimizer(cfg, model):
    if cfg["optimizer"].lower() == "sgd":
        return optim.SGD(model.parameters(), lr=cfg["lr"], momentum=cfg["momentum"])
    if cfg["optimizer"].lower() == "adam":
        return optim.Adam(
            model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"]
        )
    elif cfg["optimizer"].lower() == "adamw":
        return optim.AdamW(
            model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"]
        )


def get_scheduler(cfg, optimizer):
    if cfg["scheduler"].lower() == "steplr":
        return lr_scheduler.StepLR(
            optimizer, step_size=cfg["step_size"], gamma=cfg["gamma"]
        )
    elif cfg["scheduler"].lower() == "cosineannealinglr":
        return lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=cfg["epochs"], eta_min=cfg["eta_min"]
        )
    elif cfg["scheduler"].lower() == "reducelronplateau":
        return lr_scheduler.ReduceLROnPlateau(
            optimizer, mode="min", factor=0.5, patience=10
        )


def get_plot_fn(cfg):
    if cfg["data_name"].lower() == "darcy":
        return partial(plot_NS_2D_trainer, cmap="winter")
