"""
This module contains the Classifier class for making predictions using a trained model.
"""

import json
import torch
from torch.utils.data import DataLoader
from torch import nn
import numpy as np
from tqdm import tqdm
from icecream import ic
from .inference_dataset import InferenceDataset
from .postprocess import concat_pred_dicts, make_smooth_preds, find_true_pred_sequences

sample_config = {
    # mean and std for normalization
    "mean": (0.45211223, 0.27139644, 0.19264949),
    "std": (0.31418097, 0.21088019, 0.16059452),
    # Image Size
    "size_x": 716,
    "size_y": 716,
    # how to wrangle axes of the image before putting them in the network
    "axes": [2, 0, 1],  # 2,1,0 for opencv
    "batchsize": 16,
    "num_workers": 0,  # always 1 for Windows systems # FIXME: fix celery crash if multiprocessing
    # maybe add sigmoid after prediction?
    "activation": nn.Sigmoid(),
    "labels": [
        "appendix",
        "blood",
        "diverticule",
        "grasper",
        "ileocaecalvalve",
        "ileum",
        "low_quality",
        "nbi",
        "needle",
        "outside",
        "polyp",
        "snare",
        "water_jet",
        "wound",
    ],
}


class Classifier:
    def __init__(self, model=None, config=None, verbose=False):
        if config is None:
            config = sample_config.copy()
        self.config = config
        self.model = model
        self.verbose = verbose

    def pipe(self, paths, crops, verbose=None):
        """
        Processes input data through the model pipeline and returns predictions.
        Args:
            paths (list): List of file paths to the input data.
            crops (list): List of crop regions for the input data.
            verbose (bool, optional): If True, prints detailed logs. Defaults to None.
        Returns:
            list: Predictions generated by the model.
        """

        if verbose is None:
            verbose = self.verbose

        dataset = InferenceDataset(paths, crops, self.config)
        if verbose:
            ic("Dataset created")

        dl = DataLoader(
            dataset=dataset,
            batch_size=self.config["batchsize"],
            num_workers=self.config["num_workers"],
            shuffle=False,
            pin_memory=True,
        )
        if verbose:
            ic("Dataloader created")

        predictions = []

        with torch.inference_mode():
            if self.verbose:
                ic("Starting inference")
                
            # Ensure model exists
            if self.model is None:
                raise ValueError("Model is not loaded")
            
            # Use the device the model is currently on, with fallback to CPU
            try:
                # Check what device the model parameters are on
                model_device = next(self.model.parameters()).device
                device = model_device
                if verbose:
                    print(f"Using device: {device}")
            except StopIteration:
                # Model has no parameters, default to CPU
                device = torch.device("cpu")
                if verbose:
                    print("Model has no parameters, defaulting to CPU")
            except Exception as e:
                # Any other issue, fall back to CPU
                device = torch.device("cpu")
                if verbose:
                    print(f"Device detection failed, using CPU: {e}")
                    
            # Ensure model is in eval mode
            self.model.eval()
            
            for batch in tqdm(dl):
                batch = batch.to(device, non_blocking=True)
                prediction = self.model(batch)
                prediction = (
                    self.config["activation"](prediction).cpu().tolist()
                )  # .numpy().tolist()
                predictions += prediction

        return predictions

    def __call__(self, image, crop=None):
        return self.pipe([image], [crop])

    def readable(self, predictions):
        """
        Converts a list of predictions into a readable dictionary format.
        Args:
            predictions (list): A list of prediction values.
        Returns:
            dict: A dictionary where the keys are labels from the configuration
                  and the values are the corresponding predictions. If a prediction
                  value is of type numpy.int64, it is converted to a standard int.
        """
        readable_dict = {
            label: prediction
            for label, prediction in zip(self.config["labels"], predictions)
        }

        # when dumping to json we need to convert numpy.int64 to int
        readable_dict = {
            key: int(value) if isinstance(value, np.int64) else value
            for key, value in readable_dict.items()
        }

        return readable_dict

    def get_prediction_dict(self, predictions, paths):
        """
        Constructs a dictionary containing prediction results.
        Args:
            predictions (list): A list of prediction results.
            paths (list): A list of paths corresponding to the predictions.
        Returns:
            dict: A dictionary with the following keys:
                - "labels": The labels from the configuration.
                - "paths": The provided paths.
                - "predictions": The provided predictions.
        """

        json_dict = {
            "labels": self.config["labels"],
            "paths": paths,
            "predictions": predictions,
        }

        return json_dict

    def get_prediction_json(self, predictions, paths, json_target_path: str = None):
        """
        Saves predictions to a JSON file.

        Args:
            predictions (list): A list of prediction results.
            paths (list): A list of paths corresponding to the predictions.
            json_target_path (str, optional): The path to save the JSON file. Defaults to None.
        """
        if not json_target_path:
            json_target_path = "predictions.json"

        json_dict = self.get_prediction_dict(predictions, paths)

        with open(json_target_path, "w", encoding="utf-8") as f:
            json.dump(json_dict, f)

        if self.verbose:
            ic(f"Saved predictions to {json_target_path}")

    def post_process_predictions(
        self, pred_dicts, window_size_s=1, fps=50, min_seq_len_s=0.5
    ):
        """
        pred_dicts: list of dictionaries with the same keys
        window_size_s: size of the window in seconds for smoothing
        fps: frames per second
        min_seq_len_s: minimum length of a sequence in seconds

        Returns:
        predictions: concatenated predictions
        smooth_predictions: smoothed predictions
        binary_predictions: binary predictions
        raw_sequences: raw sequences
        filtered_sequences: filtered sequences
        """
        # Concatenate the predictions
        predictions = concat_pred_dicts(pred_dicts)

        smooth_predictions = {key: [] for key in predictions.keys()}
        for key in predictions.keys():
            smooth_predictions[key] = make_smooth_preds(
                predictions[key], window_size_s=window_size_s, fps=fps
            )

        binary_predictions = {}
        for key in smooth_predictions.keys():
            binary_predictions[key] = np.array(
                [p > 0.5 for p in smooth_predictions[key]]
            )

        raw_sequences = {}
        for key, value in binary_predictions.items():
            raw_sequences[key] = find_true_pred_sequences(value)

        filtered_sequences = {}
        min_seq_len = int(min_seq_len_s * fps)
        for key, sequences in raw_sequences.items():
            filtered_sequences[key] = [
                s for s in sequences if s[1] - s[0] > min_seq_len
            ]

        return (
            predictions,
            smooth_predictions,
            binary_predictions,
            raw_sequences,
            filtered_sequences,
        )

    def post_process_predictions_serializable(
        self, pred_dicts, window_size_s=1, fps=50, min_seq_len_s=0.5
    ):
        """
        Post-processes prediction dictionaries to make them serializable.
        This method takes prediction dictionaries, processes them to ensure all
        elements are serializable (e.g., converting numpy arrays to lists), and
        organizes the results into a dictionary of dictionaries.
        Args:
            pred_dicts (list): List of prediction dictionaries.
            window_size_s (int, optional): Window size in seconds. Defaults to 1.
            fps (int, optional): Frames per second. Defaults to 50.
            min_seq_len_s (float, optional): Minimum sequence length in seconds. Defaults to 0.5.
        Returns:
            dict: A dictionary containing processed prediction results with the following keys:
                - "predictions": Processed predictions.
                - "smooth_predictions": Smoothed predictions.
                - "binary_predictions": Binary predictions.
                - "raw_sequences": Raw sequences.
                - "filtered_sequences": Filtered sequences.
        """

        result = self.post_process_predictions(
            pred_dicts, window_size_s, fps, min_seq_len_s
        )

        for i, _dict in enumerate(result):
            _keys = list(_dict.keys())
            for key in _keys:
                # if numpy array
                if hasattr(_dict[key], "tolist"):
                    result[i][key] = _dict[key].tolist()

                # check if list of tuples
                # if so, make sure each tuple has 2 elements and split to two lists (start, stop)
                if all(isinstance(x, tuple) for x in _dict[key]):
                    if all(len(x) == 2 for x in _dict[key]):
                        result[i][f"{key}_start"] = [int(x[0]) for x in _dict[key]]
                        result[i][f"{key}_stop"] = [int(x[1]) for x in _dict[key]]
                        del result[i][key]

        # make dict of dicts
        result_dict = {
            "predictions": result[0],
            "smooth_predictions": result[1],
            "binary_predictions": result[2],
            "raw_sequences": result[3],
            "filtered_sequences": result[4],
        }

        return result_dict
