import json
import logging
from pathlib import Path
from typing import Callable, Optional

import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from torchvision.datasets import VisionDataset

from base.torchvision_dataset import TorchvisionDataset


class SubTer_Dataset(TorchvisionDataset):
    """
    Wrapper for SubTerTraining and SubTerInference, sets up train/test/inference/data_set as needed.
    """

    def __init__(
        self,
        root: str,
        num_known_normal: int = 0,
        num_known_outlier: int = 0,
        k_fold_num: int = None,
        inference: bool = False,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        seed: int = 0,
        split: float = 0.7,
    ):
        super().__init__(root, k_fold_number=k_fold_num)
        self.inference_set = None
        self.train_set = None
        self.test_set = None
        self.data_set = None

        if inference:
            self.inference_set = SubTerInference(
                root=root,
                transform=transform,
            )
            return

        # Always require the manual label file
        manual_json_path = Path(root) / "manually_labeled_anomaly_frames.json"
        if not manual_json_path.exists():
            raise FileNotFoundError(f"Required file not found: {manual_json_path}")

        # For k_fold, data_set is the full dataset, train/test are None
        if k_fold_num is not None:
            self.data_set = SubTerTraining(
                root=root,
                num_known_normal=num_known_normal,
                num_known_outlier=num_known_outlier,
                transform=transform,
                target_transform=target_transform,
                seed=seed,
                split=1.0,  # use all data for k-fold
            )
            self.train_set = None
            self.test_set = None
        else:
            # Standard split
            self.train_set = SubTerTraining(
                root=root,
                num_known_normal=num_known_normal,
                num_known_outlier=num_known_outlier,
                transform=transform,
                target_transform=target_transform,
                seed=seed,
                split=split,
                train=True,
            )
            self.test_set = SubTerTraining(
                root=root,
                num_known_normal=num_known_normal,
                num_known_outlier=num_known_outlier,
                transform=transform,
                target_transform=target_transform,
                seed=seed,
                split=split,
                train=False,
            )
            self.data_set = None  # not used unless k_fold

    def get_file_name_from_idx(self, idx: int) -> Optional[str]:
        """
        Get filename for a file_id by delegating to the appropriate dataset.

        Args:
            idx: The file index to look up

        Returns:
            str: The filename corresponding to the index, or None if not found
        """

        # For non-inference, use any available dataset (they all have the same files)
        if self.data_set is not None:
            return self.data_set.get_file_name_from_idx(idx)
        if self.train_set is not None:
            return self.train_set.get_file_name_from_idx(idx)
        if self.test_set is not None:
            return self.test_set.get_file_name_from_idx(idx)
        if self.inference_set is not None:
            return self.inference_set.get_file_name_from_idx(idx)

        return None


class SubTerTraining(VisionDataset):
    """
    Loads all data, builds targets, and supports train/test split.
    """

    def __init__(
        self,
        root: str,
        num_known_normal: int = 0,
        num_known_outlier: int = 0,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        seed: int = 0,
        split: float = 0.7,
        train: bool = True,
    ):
        super().__init__(root, transform=transform, target_transform=target_transform)
        logger = logging.getLogger()

        manual_json_path = Path(root) / "manually_labeled_anomaly_frames.json"
        with open(manual_json_path, "r") as f:
            manual_data = json.load(f)
            manual_anomaly_ranges = {
                item["filename"]: (
                    item["semi_target_begin_frame"],
                    item["semi_target_end_frame"],
                )
                for item in manual_data["files"]
            }

        all_data = []
        all_file_ids = []
        all_frame_ids = []
        all_filenames = []
        test_target_experiment_based = []
        test_target_manually_set = []
        train_semi_targets = []
        file_names = {}
        file_idx = 0

        for experiment_file in sorted(Path(root).iterdir()):
            if experiment_file.suffix != ".npy":
                continue

            file_names[file_idx] = experiment_file.name
            experiment_data = np.load(experiment_file)
            n_frames = experiment_data.shape[0]

            is_smoke = "smoke" in experiment_file.name
            if is_smoke:
                if experiment_file.name not in manual_anomaly_ranges:
                    raise ValueError(
                        f"Experiment file {experiment_file.name} is marked as smoke but has no manual anomaly ranges."
                    )
                manual_anomaly_start_frame, manual_anomaly_end_frame = (
                    manual_anomaly_ranges[experiment_file.name]
                )

            # Experiment-based: 1 (normal), -1 (anomaly)
            exp_based_targets = (
                np.full(n_frames, -1, dtype=np.int8)  # anomaly
                if is_smoke
                else np.full(n_frames, 1, dtype=np.int8)  # normal
            )

            # Manually set: 1 (normal), -1 (anomaly), 0 (unknown/NaN)
            if not is_smoke:
                manual_targets = np.full(n_frames, 1, dtype=np.int8)  # normal
            else:
                manual_targets = np.zeros(n_frames, dtype=np.int8)  # unknown
                manual_targets[
                    manual_anomaly_start_frame:manual_anomaly_end_frame
                ] = -1  # anomaly

            # log how many manual anomaly frames were set to each value
            logger.info(
                f"Experiment {experiment_file.name}: "
                f"Manual targets - normal(1): {np.sum(manual_targets == 1)}, "
                f"anomaly(-1): {np.sum(manual_targets == -1)}, "
                f"unknown(0): {np.sum(manual_targets == 0)}"
            )

            # Semi-supervised targets: 1 (known normal), -1 (known anomaly), 0 (unknown)
            if not is_smoke:
                semi_targets = np.ones(n_frames, dtype=np.int8)  # normal
            else:
                semi_targets = np.zeros(n_frames, dtype=np.int8)  # unknown
                semi_targets[
                    manual_anomaly_start_frame:manual_anomaly_end_frame
                ] = -1  # anomaly

            all_data.append(experiment_data)
            all_file_ids.append(np.full(n_frames, file_idx, dtype=np.int32))
            all_frame_ids.append(np.arange(n_frames, dtype=np.int32))
            all_filenames.extend([experiment_file.name] * n_frames)
            test_target_experiment_based.append(exp_based_targets)
            test_target_manually_set.append(manual_targets)
            train_semi_targets.append(semi_targets)

            file_idx += 1

        # Flatten everything
        data = np.nan_to_num(np.concatenate(all_data))
        file_ids = np.concatenate(all_file_ids)
        frame_ids = np.concatenate(all_frame_ids)
        filenames = all_filenames
        self.file_names = file_names

        test_target_experiment_based = np.concatenate(test_target_experiment_based)
        test_target_manually_set = np.concatenate(test_target_manually_set)
        semi_targets_np = np.concatenate(train_semi_targets)

        # Limit the number of known normal/anomaly samples for training
        np.random.seed(seed)
        normal_indices = np.where(semi_targets_np == 1)[0]
        anomaly_indices = np.where(semi_targets_np == -1)[0]

        if num_known_normal > 0 and len(normal_indices) > num_known_normal:
            keep_normal = np.random.choice(
                normal_indices, size=num_known_normal, replace=False
            )
        else:
            keep_normal = normal_indices

        if num_known_outlier > 0 and len(anomaly_indices) > num_known_outlier:
            keep_anomaly = np.random.choice(
                anomaly_indices, size=num_known_outlier, replace=False
            )
        else:
            keep_anomaly = anomaly_indices

        semi_targets_np[(semi_targets_np == 1) | (semi_targets_np == -1)] = 0
        semi_targets_np[keep_normal] = 1
        semi_targets_np[keep_anomaly] = -1

        # Shuffle and split
        indices = np.arange(len(data))
        np.random.seed(seed)
        np.random.shuffle(indices)
        split_idx = int(split * len(data))
        if train:
            use_idx = indices[:split_idx]
        else:
            use_idx = indices[split_idx:]

        self.data = torch.tensor(data[use_idx])
        self.file_ids = file_ids[use_idx]
        self.frame_ids = frame_ids[use_idx]
        self.filenames = [filenames[i] for i in use_idx]
        self.test_target_experiment_based = torch.tensor(
            test_target_experiment_based[use_idx], dtype=torch.int8
        )
        self.test_target_manually_set = torch.tensor(
            test_target_manually_set[use_idx], dtype=torch.int8
        )

        # log how many of the test_target_manually_set are in each category
        logger.info(
            f"Test targets - normal(1): {np.sum(self.test_target_manually_set.numpy() == 1)}, "
            f"anomaly(-1): {np.sum(self.test_target_manually_set.numpy() == -1)}, "
            f"unknown(0): {np.sum(self.test_target_manually_set.numpy() == 0)}"
        )

        self.train_semi_targets = torch.tensor(
            semi_targets_np[use_idx], dtype=torch.int8
        )

        self.transform = transform if transform else transforms.ToTensor()
        self.target_transform = target_transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img = self.data[index]
        target_experiment_based = int(self.test_target_experiment_based[index])
        target_manually_set = int(self.test_target_manually_set[index])
        semi_target = int(self.train_semi_targets[index])
        file_id = int(self.file_ids[index])
        frame_id = int(self.frame_ids[index])

        img = Image.fromarray(img.numpy(), mode="F")
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target_experiment_based = self.target_transform(target_experiment_based)
            target_manually_set = self.target_transform(target_manually_set)
            semi_target = self.target_transform(semi_target)

        return (
            img,
            target_experiment_based,
            target_manually_set,
            semi_target,
            index,
            (file_id, frame_id),
        )

    def get_file_name_from_idx(self, idx: int):
        return self.file_names.get(idx, None)


class SubTerInference(VisionDataset):
    """
    Loads a single experiment file for inference.
    """

    def __init__(
        self,
        root: str,
        transform: Optional[Callable] = None,
    ):
        super().__init__(root, transform=transform)
        logger = logging.getLogger()

        experiment_file = Path(root)
        if not experiment_file.is_file():
            logger.error(
                "For inference the data path has to be a single experiment file!"
            )
            raise Exception("Inference data is not a loadable file!")

        self.data = np.load(experiment_file)
        self.data = np.nan_to_num(self.data)
        self.data = torch.tensor(self.data)
        self.filenames = [experiment_file.name] * self.data.shape[0]
        self.file_ids = np.zeros(self.data.shape[0], dtype=np.int32)
        self.frame_ids = np.arange(self.data.shape[0], dtype=np.int32)
        self.file_names = {0: experiment_file.name}

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img = self.data[index]
        file_id = int(self.file_ids[index])
        frame_id = int(self.frame_ids[index])

        img = Image.fromarray(img.numpy(), mode="F")
        if self.transform is not None:
            img = self.transform(img)

        return img, index, (file_id, frame_id)

    def get_file_name_from_idx(self, idx: int):
        return self.file_names.get(idx, None)
