import argparse
import os.path
import pickle
from typing import Any, Optional

import numpy as np
import torch.nn
from torch import inference_mode

from datasets.numpy import NumpyDataset, FolderNumpyDataset
from src.datasets.noise import NumpyNoiseDataset
from src.datasets.noise_label import NoiseLabelDataset, NumpyNoiseLabelDataset
from src.fid.utils import calculate_fid_for_dataset_and_model, RandomNoiseDataset, RandomLabelDataset
from src.models.models.edm import create_edm_model
from src.utils.load import load_from_state_dict
from src.utils.utils import create_one_hot_torch
from torch_utils.distributed.distributed_manager import DistributedManager
from torch_utils.utils import get_default_device
from utils.logger.logger import Logger
from utils.utils import get_class_name


@inference_mode()
def apply_model(model: torch.nn.Module, x: torch.Tensor, sigma: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
    return model(x, sigma, label)


@inference_mode()
def inference_single_step(
        model: torch.nn.Module,
        noises: torch.Tensor,
        labels: torch.Tensor = None,
        num_classes: Optional[int] = None,
        multiply_noises: bool = True,
        sigma: float = 80.0
) -> torch.Tensor:
    assert labels is None or num_classes is not None, 'labels must be None or num_classes must not be None'
    batch_size: int = noises.shape[0]
    return model(
        (noises * sigma if multiply_noises else noises).to(get_default_device()),
        torch.ones(batch_size).to(get_default_device()) * sigma,
        create_one_hot_torch(labels, num_classes).to(get_default_device()) if labels is not None else None
    )


def run(
        image_height: int,
        image_width: int,
        image_channels: int,
        reference_path: str,
        model_name: str,
        model_path: str,
        noise_folder: str = None,
        label_folder: str = None,
        noise_start_index: int = 0,
        label_start_index: int = 0,
        num_classes: Optional[int] = None,
        model_load_keys: list[str] = None,
        num_samples: int = 50_000,
        batch_size: int = 50,
        output_save_path: str = None,
        sigma: float = 80.0,
        dims: int = 2048,
        data_loader_workers: int = 8,
        distributed: bool = False,
        distributed_backend: str = 'nccl'
) -> None:
    Logger.debug(
        f'{get_class_name(run)} - '
        f'image_height: {image_height}, '
        f'image_width: {image_width}, '
        f'image_channels: {image_channels}, '
        f'reference_path: {reference_path}, '
        f'model_name: {model_name}, '
        f'model_path: {model_path}, '
        f'noise_folder: {noise_folder}, '
        f'label_folder: {label_folder}, '
        f'noise_start_index: {noise_start_index}, '
        f'label_start_index: {label_start_index}, '
        f'num_classes: {num_classes}, '
        f'model_load_keys: {model_load_keys}, '
        f'num_samples: {num_samples}, '
        f'batch_size: {batch_size}, '
        f'output_save_path: {output_save_path}, '
        f'sigma: {sigma}, '
        f'dims: {dims}, '
        f'data_loader_workers: {data_loader_workers}, '
        f'distributed: {distributed}, '
        f'distributed_backend: {distributed_backend}'
    )
    if distributed:
        DistributedManager.init(backend=distributed_backend)

    conditional: bool = num_classes is not None
    Logger.debug(f'conditional: {conditional}')

    Logger.debug('creating model')
    model: torch.nn.Module = create_edm_model(model_name)
    Logger.debug('loading model')
    model: torch.nn.Module = load_from_state_dict(model, model_path, model_load_keys)
    model.to(get_default_device())
    model.eval()

    noise_dataset: NumpyDataset = FolderNumpyDataset(
        data_shape=(image_channels, image_height, image_width),
        data_type=np.float64,
        folder=noise_folder,
        num_samples=num_samples,
        start_index=noise_start_index
    ) if noise_folder is not None else RandomNoiseDataset(
        num_samples=num_samples,
        noise_shape=(image_channels, image_height, image_width)
    )
    label_dataset: NumpyDataset = FolderNumpyDataset(
        data_shape=(),
        data_type=np.int64,
        folder=label_folder,
        num_samples=num_samples,
        start_index=label_start_index
    ) if label_folder is not None else RandomLabelDataset(
        num_samples=num_samples,
        num_classes=num_classes
    )

    dataset: NoiseLabelDataset = NumpyNoiseLabelDataset(
        noise_dataset=noise_dataset,
        label_dataset=label_dataset
    ) if conditional else NumpyNoiseDataset(
        noise_dataset=noise_dataset
    )

    Logger.debug('calculating fid')
    fid: float = calculate_fid_for_dataset_and_model(
        model=model,
        dataset=dataset,
        reference_path=reference_path,
        inference_batch_func=lambda m, n, l: inference_single_step(
            model=m,
            noises=n,
            labels=l if conditional else None,
            num_classes=num_classes,
            sigma=sigma
        ),
        batch_size=batch_size,
        dims=dims,
        data_loader_workers=data_loader_workers
    )
    Logger.debug(f'fid: {fid}')
    if DistributedManager.is_main():
        if output_save_path is not None:
            Logger.debug(f'saving fid to {output_save_path}')
            if os.path.dirname(output_save_path) != '':
                os.makedirs(os.path.dirname(output_save_path), exist_ok=True)
            with open(output_save_path, 'wb') as file:
                pickle.dump(fid, file)

    if distributed:
        DistributedManager.destroy()


def run_from_config(config: dict[str, Any]) -> None:
    run(**config)


def parse_args() -> argparse.Namespace:
    parser: argparse.ArgumentParser = argparse.ArgumentParser()
    parser.add_argument('--image_height', type=int, required=True)
    parser.add_argument('--image_width', type=int, required=True)
    parser.add_argument('--image_channels', type=int, required=True)
    parser.add_argument('--reference_path', type=str, required=True)
    parser.add_argument('--model_name', type=str, required=True)
    parser.add_argument('--model_path', type=str, required=True)
    parser.add_argument('--noise_folder', type=str, default=None)
    parser.add_argument('--label_folder', type=str, default=None)
    parser.add_argument('--noise_start_index', type=int, default=0)
    parser.add_argument('--label_start_index', type=int, default=0)
    parser.add_argument('--num_classes', type=int, default=None)
    parser.add_argument('--model_load_keys', type=str, nargs='+', default=None)
    parser.add_argument('--num_samples', type=int, default=50_000)
    parser.add_argument('--batch_size', type=int, default=50)
    parser.add_argument('--output_save_path', type=str, default=None)
    parser.add_argument('--sigma', type=float, default=80.0)
    parser.add_argument('--dims', type=int, default=2048)
    parser.add_argument('--data_loader_workers', type=int, default=8)
    parser.add_argument('--distributed', action='store_true')
    parser.add_argument('--distributed_backend', type=str, default='nccl')
    return parser.parse_args()


def get_config_from_args(args: argparse.Namespace) -> dict[str, Any]:
    return vars(args)


def main() -> None:
    args: argparse.Namespace = parse_args()
    config: dict[str, Any] = get_config_from_args(args)
    run_from_config(config)


if __name__ == '__main__':
    main()
