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

import numpy as np
import torch
from tqdm import tqdm

from fid_pytorch.fid_pytorch import create_inception_model, get_activations, \
    calculate_statistics_from_activations, load_fid_stats, calculate_frechet_distance
from fid_pytorch.scripts.utils import NumpyArrayDataset
from src.fid_utils.inference.multi_step_euler import inference_multi_step_euler_batch
from src.models.models.edm import create_edm_model
from src.utils.load import load_from_state_dict
from utils.logger.logger import Logger
from utils.numpy.load import load
from utils.utils import get_object_name, get_class_name


def fid_multi_step_euler(
        num_steps: int,
        time_steps: list[int],
        model: torch.nn.Module,
        image_height: int,
        image_width: int,
        image_channels: int,
        reference_path: str,
        num_classes: Optional[int] = None,
        num_samples: int = 50_000,
        batch_size: int = 50,
        noises_folder: str = None,
        labels_folder: str = None,
        start_noise_index: int = 0,
        start_label_index: int = 0,
        randn_like: Callable[[np.ndarray], np.ndarray] = lambda x: np.random.randn(*x.shape),
        sigma_min: float = 0.002,
        sigma_max: float = 80,
        rho: float = 7,
        s_churn: float = 0,
        s_min: float = 0,
        s_max: float = float('inf'),
        s_noise: float = 1,
        multiply_noises: bool = True,
        n_processes: int = 8,
        dims: int = 2048,
        device: str = None
) -> float:
    print(
        f'{get_class_name(fid_multi_step_euler)} - '
        f'num_steps: {num_steps}, '
        f'time_steps: {time_steps}, '
        f'model: {get_object_name(model)}, '
        f'image_height: {image_height}, '
        f'image_width: {image_width}, '
        f'image_channels: {image_channels}, '
        f'reference_path: {reference_path}, '
        f'num_classes: {num_classes}, '
        f'num_samples: {num_samples}, '
        f'batch_size: {batch_size}, '
        f'noises_folder: {noises_folder}, '
        f'labels_folder: {labels_folder}, '
        f'start_noise_index: {start_noise_index}, '
        f'start_label_index: {start_label_index}, '
        f'sigma_min: {sigma_min}, '
        f'sigma_max: {sigma_max}, '
        f'rho: {rho}, '
        f's_churn: {s_churn}, '
        f's_min: {s_min}, '
        f's_max: {s_max}, '
        f's_noise: {s_noise}, '
        f'multiply_noises: {multiply_noises}, '
        f'n_processes: {n_processes}, '
        f'dims: {dims}, '
        f'device: {device}'
    )
    n_batches: int = int(np.ceil(num_samples / batch_size))
    outputs: np.ndarray = np.zeros((num_samples, image_channels, image_height, image_width))
    for batch_index in tqdm(range(n_batches)):
        start_index = batch_index * batch_size
        end_index = min(start_index + batch_size, num_samples)
        noises_batch: np.ndarray = load(
            folder=noises_folder,
            n_samples=end_index - start_index,
            start_index=start_index + start_noise_index,
            n_processes=n_processes
        ) if noises_folder is not None else (
            np.random.randn(end_index - start_index, image_channels, image_height, image_width)
        )
        labels_batch: np.ndarray = (load(
            folder=labels_folder,
            n_samples=end_index - start_index,
            start_index=start_index + start_label_index,
            n_processes=n_processes
        ) if labels_folder is not None else (
            np.random.randint(0, num_classes, end_index - start_index)
        )) if num_classes is not None else None
        outputs[start_index:end_index] = inference_multi_step_euler_batch(
            num_steps=num_steps,
            time_steps=time_steps,
            model=model,
            noises=noises_batch,
            num_classes=num_classes,
            labels=labels_batch,
            randn_like=randn_like,
            sigma_min=sigma_min,
            sigma_max=sigma_max,
            rho=rho,
            s_churn=s_churn,
            s_min=s_min,
            s_max=s_max,
            s_noise=s_noise,
            multiply_noises=multiply_noises,
            device=device
        )

    dataset: NumpyArrayDataset = NumpyArrayDataset(outputs)

    activations: np.ndarray = get_activations(
        dataset,
        model=create_inception_model(dims).to(device),
        batch_size=batch_size,
        dims=dims,
        device=device,
        num_workers=n_processes
    )
    mu, sigma = calculate_statistics_from_activations(activations)
    mu_ref, sigma_ref = load_fid_stats(reference_path)
    fid: float = calculate_frechet_distance(mu, sigma, mu_ref, sigma_ref)
    Logger.debug(f'fid: {fid}')
    return fid


def run(
        num_steps: int,
        time_steps: list[int],
        model_name: str,
        model_load_path: str,
        image_height: int,
        image_width: int,
        image_channels: int,
        reference_path: str,
        num_classes: Optional[int] = None,
        num_samples: int = 50_000,
        batch_size: int = 50,
        noises_folder: str = None,
        labels_folder: str = None,
        save_path: str = None,
        start_noise_index: int = 0,
        start_label_index: int = 0,
        sigma_min: float = 0.002,
        sigma_max: float = 80,
        rho: float = 7,
        s_churn: float = 0,
        s_min: float = 0,
        s_max: float = float('inf'),
        s_noise: float = 1,
        model_load_keys: list[str] = None,
        multiply_noises: bool = True,
        n_processes: int = 8,
        device: str = None
) -> None:
    print(
        f'{get_class_name(run)} - '
        f'num_steps: {num_steps}, '
        f'time_steps: {time_steps}, '
        f'model_name: {model_name}, '
        f'model_load_path: {model_load_path}, '
        f'image_height: {image_height}, '
        f'image_width: {image_width}, '
        f'image_channels: {image_channels}, '
        f'reference_path: {reference_path}, '
        f'num_classes: {num_classes}, '
        f'num_samples: {num_samples}, '
        f'batch_size: {batch_size}, '
        f'noises_folder: {noises_folder}, '
        f'labels_folder: {labels_folder}, '
        f'start_noise_index: {start_noise_index}, '
        f'start_label_index: {start_label_index}, '
        f'sigma_min: {sigma_min}, '
        f'sigma_max: {sigma_max}, '
        f'rho: {rho}, '
        f's_churn: {s_churn}, '
        f's_min: {s_min}, '
        f's_max: {s_max}, '
        f's_noise: {s_noise}, '
        f'model_load_keys: {model_load_keys}, '
        f'multiply_noises: {multiply_noises}, '
        f'n_processes: {n_processes}, '
        f'device: {device}'
    )
    Logger.debug('creating model')
    model: torch.nn.Module = create_edm_model(model_name)
    model = load_from_state_dict(model, load_path=model_load_path, load_keys=model_load_keys)
    model.eval()
    model.to(device)

    fid: float = fid_multi_step_euler(
        num_steps=num_steps,
        time_steps=time_steps,
        model=model,
        noises_folder=noises_folder,
        labels_folder=labels_folder,
        num_samples=num_samples,
        batch_size=batch_size,
        image_height=image_height,
        image_width=image_width,
        image_channels=image_channels,
        reference_path=reference_path,
        num_classes=num_classes,
        start_noise_index=start_noise_index,
        start_label_index=start_label_index,
        sigma_min=sigma_min,
        sigma_max=sigma_max,
        rho=rho,
        s_churn=s_churn,
        s_min=s_min,
        s_max=s_max,
        s_noise=s_noise,
        multiply_noises=multiply_noises,
        n_processes=n_processes,
        device=device
    )
    Logger.debug(f'fid: {fid}')

    if save_path is not None:
        Logger.debug(f'saving fid to {save_path}')
        if os.path.dirname(save_path) != '':
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
        with open(save_path, 'wb') as file:
            pickle.dump(fid, file)


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


def parse_args() -> argparse.Namespace:
    parser: argparse.ArgumentParser = argparse.ArgumentParser()
    parser.add_argument('--num_steps', type=int, required=True)
    parser.add_argument('--time_steps', type=int, nargs='+', required=True)
    parser.add_argument('--model_name', type=str, required=True)
    parser.add_argument('--model_load_path', type=str, required=True)
    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('--num_classes', type=int, default=None)
    parser.add_argument('--num_samples', type=int, default=50_000)
    parser.add_argument('--batch_size', type=int, default=50)
    parser.add_argument('--noises_folder', type=str, default=None)
    parser.add_argument('--labels_folder', type=str, default=None)
    parser.add_argument('--save_path', type=str, default=None)
    parser.add_argument('--start_noise_index', type=int, default=0)
    parser.add_argument('--start_label_index', type=int, default=0)
    parser.add_argument('--sigma_min', type=float, default=0.002)
    parser.add_argument('--sigma_max', type=float, default=80.0)
    parser.add_argument('--rho', type=float, default=7.0)
    parser.add_argument('--s_churn', type=float, default=0.0)
    parser.add_argument('--s_min', type=float, default=0.0)
    parser.add_argument('--s_max', type=float, default=float('inf'))
    parser.add_argument('--s_noise', type=float, default=1.0)
    parser.add_argument('--model_load_keys', type=str, nargs='+', default=None)
    parser.add_argument('--multiply_noises', type=bool, default=True)
    parser.add_argument('--n_processes', type=int, default=8)
    parser.add_argument('--device', type=str, default=None)
    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()
