import json
import pathlib

from modules.dataLoader import BaseDataLoader
from modules.model.BaseModel import BaseModel
from modules.modelLoader.BaseModelLoader import BaseModelLoader
from modules.modelSetup.BaseModelSetup import BaseModelSetup
from modules.util import create
from modules.util.config.TrainConfig import QuantizationConfig, TrainConfig
from modules.util.torch_util import torch_gc
from modules.util.TrainProgress import TrainProgress

import torch

from tqdm import tqdm


class GenerateLossesModel:
    """Based on train args, writes a JSON instead of a model with filenames mapped to losses,
    in order of decreasing loss."""
    config: TrainConfig
    train_device: torch.device
    temp_device: torch.device
    model_loader: BaseModelLoader
    model_setup: BaseModelSetup
    data_loader: BaseDataLoader
    model: BaseModel

    def __init__(self, config: TrainConfig, output_path: str):
        # Create a copy of args because we will mutate
        # the batch size and gradient accumulation steps.
        config = TrainConfig.default_values().from_dict(config.to_dict())
        config.batch_size = 1
        config.gradient_accumulation_steps = 1

        self.config = config
        self.output_path = output_path
        self.train_device = torch.device(self.config.train_device)
        self.temp_device = torch.device(self.config.temp_device)

    def start(self):
        output_dir = pathlib.Path(self.output_path).parent
        if not output_dir.is_dir():
            output_dir.mkdir(parents=True, exist_ok=True)

        if self.config.train_dtype.enable_tf():
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True

        self.model_loader = create.create_model_loader(self.config.model_type, self.config.training_method)
        self.model_setup = create.create_model_setup(
            self.config.model_type,
            self.train_device,
            self.temp_device,
            self.config.training_method,
            self.config.debug_mode,
        )

        model_names = self.config.model_names()

        self.model = self.model_loader.load(
            model_type=self.config.model_type,
            model_names=model_names,
            weight_dtypes=self.config.weight_dtypes(),
            quantization=QuantizationConfig.default_values(),
        )
        self.model.train_config = self.config

        self.model_setup.setup_optimizations(self.model, self.config)
        self.model_setup.setup_model(self.model, self.config)
        self.model.eval()
        self.model.train_progress = TrainProgress()
        torch_gc()

        self.data_loader = create.create_data_loader(
            train_device=self.train_device,
            temp_device=self.temp_device,
            model=self.model,
            model_type=self.config.model_type,
            model_setup=self.model_setup,
            training_method=self.config.training_method,
            config=self.config,
            train_progress=self.model.train_progress,
            is_validation=False
        )

        data_set = self.data_loader.get_data_set()
        data_set.start_next_epoch()
        approximate_num_steps = data_set.approximate_length()

        step_tqdm = tqdm(self.data_loader.get_data_loader(), total=approximate_num_steps, desc="Calculating loss")

        self.model_setup.setup_train_device(self.model, self.config)

        filename_loss_list: list[tuple[str, float]] = []
        # Don't really need a backward pass here, so we can make the calculation MUCH faster.
        with torch.inference_mode():
            for batch in step_tqdm:
                    model_output_data = self.model_setup.predict(
                        self.model,
                        batch,
                        self.config,
                        self.model.train_progress,
                        deterministic=True,
                    )
                    loss = self.model_setup.calculate_loss(
                        self.model,
                        batch,
                        model_output_data,
                        self.config,
                    )
                    filename_loss_list.append((batch['image_path'][0], float(loss)))

        # Sort such that highest loss comes first
        filename_loss_list.sort(key=lambda x: x[1], reverse=True)
        filename_to_loss: dict[str, float] = {x[0]: x[1] for x in filename_loss_list}
        with open(self.output_path, "w") as f:
            json.dump(filename_to_loss, f, indent=4)
