# indad/indad.py

import os
import sys
from pathlib import Path
import imghdr
from datetime import datetime
import random
import shutil
import json
from easydict import EasyDict

import numpy as np
from PIL import Image
import math

import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from sklearn.model_selection import KFold
from glob import glob
sys.path.insert(0, os.path.abspath(
    os.path.join(os.path.dirname(__file__), './../../')
))

from ecos_core.indad.model_builder import SPADE, PaDiM, PatchCore, \
    AmazonPatchCore, PatchCoreFAISS, SAPatchCore, PatchCoreCluster
from ecos_core.indad.utils_indad import save_image
from ecos_core.base_model import BaseModel
from ecos_core.indad.sample_code.data_transform import transform_default, get_transform
from ecos_core.common.tools.auto_augment import create_transform_function
from ecos_core.indad.sample_code.data_setup import TrainDatasetListPath
K_NEIGHBORS = 3
D_REDUCED = 350
F_CORESET = .10
N_CLUSTER = 100
BACKBONE_NAME = "wide_resnet50_2"
NORMALIZE = False
M_SPADE = f"spade"
M_PADIM = f"padim"
M_PATCHCORE = f"patchcore"
M_AMAZONE_PATCHCORE = "amazon_patchcore"
M_PATCHCORE_FAISS = "patchcore_faiss"
M_SA_PATCHCORE = "sa_patchcore"
M_PATCHCORE_CLUSTER = "patchcore_cluster"
ALLOWED_METHODS = ["spade", "padim", "patchcore", "amazon_patchcore", "patchcore_faiss", "sa_patchcore", M_PATCHCORE_CLUSTER]
# config amazon patchcore
PRETRAIN_EMBED_DIMENSION = 1024
TARGET_EMBED_DIMENSION = 1024
PATCHSIZE = 3
LAYER_EXTRACT = ["layer2", "layer3"]
ANOMALY_SCORER_NUM_NN = 1
SEED = 0
PERCENTAGE= 0.1
SIZE = 224
SECOND_SIZE = 512
NUM_WORKERS = 4
MAX_ITER = 300 # max iteration Kmeans
# Now, only support for CPU
IS_CUDA_USED = False


FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH


sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


class Model(BaseModel):
    """This is the base model for INDAD

    Args:
        BaseModel (Class): Common class for all ecos_core models

    Returns: None
    """
    NAME = "indad"
    def __init__(self, opt):
        """This function is called when creating a new instance of INDAD Model

        Args:
            opt (NameSpace): The default options neccessary for initializing the model. 
                             opt should includes these attributes: class_names, weight, device.

        """
        super().__init__(opt)
        self.opt = opt
        self.opt.class_names = ["good", "anomalous"]
        self.opt.weight = "./model_weight.pt"
        self.opt.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        self.check_init()

    def check_init(self):
        # check pixel score is used
        if not hasattr(self.opt, "is_pixel_score_used"):
            self.opt.is_pixel_score_used = False

        if not hasattr(self.opt, "display_mode"):
            self.opt.display_mode = "Heatmap"

        if hasattr(self.opt, 'model_parameters'):
            global K_NEIGHBORS, SIZE, BACKBONE_NAME, NORMALIZE, MAX_ITER, SECOND_SIZE
            K_NEIGHBORS = self.opt.model_parameters.get("k", K_NEIGHBORS)
            SIZE = self.opt.model_parameters.get("size", SIZE)
            BACKBONE_NAME = self.opt.model_parameters.get("backbone_name", BACKBONE_NAME)
            NORMALIZE = self.opt.model_parameters.get("normalize", NORMALIZE)
            MAX_ITER = self.opt.model_parameters.get("max_iter", MAX_ITER)
            SECOND_SIZE = self.opt.model_parameters.get("second_size", SECOND_SIZE)
        
        if not hasattr(self.opt, "core_version"):
            self.opt.core_version = "1.1.0"

    def get_transform(self):
        """Transform input by ImageNet normalization. 

        Returns:
            Function: A function for transform a single image from path
        """

        IMAGENET_MEAN = torch.tensor([.485, .456, .406])
        IMAGENET_STD = torch.tensor([.229, .224, .225])

        data_transforms=transforms.Compose([
            transforms.Resize((SIZE,SIZE), interpolation=transforms.InterpolationMode.BICUBIC),
            transforms.ToTensor(),
            transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
        ])

        def transform(image_path):
            raw_image = Image.open(image_path).convert('RGB')

            image_transform = data_transforms(raw_image)

            image_transform = image_transform.unsqueeze(dim=0)

            return raw_image, image_transform
        
        return transform

    @staticmethod
    def generate_instance_model(method):
        """generate instance model

        Args:
            method (str): spade or padim or patchcore or amazon_patchcore method

        Returns:
            object: instance of method
        """
        if method == M_SPADE:
            instance = SPADE(
                k=K_NEIGHBORS,
                size=SIZE,
                backbone_name=BACKBONE_NAME,
            )
        elif method == M_PADIM:
            instance = PaDiM(
                size=SIZE,
                d_reduced=D_REDUCED,
                backbone_name=BACKBONE_NAME,
            )
        elif method == M_PATCHCORE:
            instance = PatchCore(
                k=K_NEIGHBORS,
                size=SIZE,
                f_coreset=F_CORESET, 
                backbone_name=BACKBONE_NAME,
                layers=LAYER_EXTRACT,
                normalize=NORMALIZE
            )
        elif method == M_AMAZONE_PATCHCORE:
            instance = AmazonPatchCore(
                device="cpu",
                backbone_name=BACKBONE_NAME,
                layers_to_extract_from=LAYER_EXTRACT,
                pretrain_embed_dimension=PRETRAIN_EMBED_DIMENSION,
                target_embed_dimension=TARGET_EMBED_DIMENSION,
                patchsize=PATCHSIZE,
                input_shape=(3, SIZE, SIZE),
                percentage=PERCENTAGE,
                anomaly_scorer_num_nn=ANOMALY_SCORER_NUM_NN,
                seed=0,
                num_workers=NUM_WORKERS,
                is_cuda_used=IS_CUDA_USED
            )
        elif method == M_PATCHCORE_FAISS:
            instance = PatchCoreFAISS(
                k=K_NEIGHBORS,
                size=SIZE,
                f_coreset=F_CORESET, 
                backbone_name=BACKBONE_NAME,
                normalize=NORMALIZE
            )
        elif method == M_SA_PATCHCORE:
            instance = SAPatchCore(
                k=K_NEIGHBORS,
                size=SIZE,
                f_coreset=F_CORESET, 
                backbone_name=BACKBONE_NAME,
                layers=LAYER_EXTRACT,
                normalize=NORMALIZE
            )
        elif method == M_PATCHCORE_CLUSTER:
            instance = PatchCoreCluster(
                size=SIZE,
                second_size=SECOND_SIZE,
                n_cluster = N_CLUSTER,
                backbone_name=BACKBONE_NAME,
                layers=LAYER_EXTRACT,
                normalize=NORMALIZE,
                max_iter=MAX_ITER
            )

        return instance

    def build_model(self, custom_weight:str=None):
        """Build model in ensemble model.
           Type of model is defined in opt.json
        
        Args:
            custom_weight(str): weight to load that not defined in opt. Default: None.
        
        Returns:
            (torch.Model): Model after built
        
        Examples:
            >>> from argparse import ArgumentParser
            >>> from ecos_core.indad.indad import Model
            >>> parser = ArgumentParser()
            >>> opts = parser.parse_args()
            >>> with open(opt_path, 'r') as f:
            >>>     opts.__dict__ = json.load(f)
            >>> model = Model(opts)
            >>> #Build model with default options
            >>> model.build_model()
            >>> #Build model with custom weight
            >>> model.build_model(custom_weight=<PATH_TO_WEIGHT>)
        """
        self.check_init()
        method = self.opt.method
        print(method)
        print(BACKBONE_NAME)
        if custom_weight is None:
            weight_path = os.path.join(str(ROOT), self.opt.weight)
        else:
            weight_path = custom_weight

        assert method in ALLOWED_METHODS, f"Select from {ALLOWED_METHODS}."
        self.opt.device = "cpu"
        checkpoint = torch.load(weight_path, map_location=self.opt.device)

        model = self.generate_instance_model(method)

        if method == M_SPADE:
            model.load_state_dict(checkpoint['model_state_dict'])
            model.z_lib = checkpoint['z_lib']
            model.feature_maps = checkpoint['features_maps']
        elif method == M_PADIM:
            model.load_state_dict(checkpoint['model_state_dict'])
            model.resize = checkpoint['resize']
            model.r_indices = checkpoint['r_indices']
            model.means_reduced = checkpoint['means_reduced']
            model.E_inv = checkpoint['E_inv']
        elif method == M_PATCHCORE or method == M_SA_PATCHCORE:
            model.load_state_dict(checkpoint['model_state_dict'])
            model.resize = checkpoint['resize']
            model.patch_lib = checkpoint['patch_lib']
        elif method == M_AMAZONE_PATCHCORE:
            patchcore_params = checkpoint['patchcore_params']
            patchcore_params["backbone"] = AmazonPatchCore.get_backbone_from_name(
                patchcore_params["backbone.name"]
            )
            patchcore_params["backbone"].name = patchcore_params["backbone.name"]
            model.load(**patchcore_params, device=torch.device(self.opt.device), nn_method=model.nn_method)
            model.anomaly_scorer.nn_method.search_index = checkpoint["anomaly_scorer"]
        elif method == M_PATCHCORE_FAISS:
            model.load_state_dict(checkpoint['model_state_dict'])
            model.resize = checkpoint['resize']
            model.patch_lib = checkpoint['patch_lib']
            model.anomaly_scorer.nn_method.search_index = checkpoint.get('anomaly_scorer_index',None)
        elif method == M_PATCHCORE_CLUSTER:
            model.load_state_dict(checkpoint['model_state_dict'])
            model.resize = checkpoint['resize']
            model.second_resize = checkpoint['second_resize']
            model.patch_lib = checkpoint['patch_lib']
            model.kmeans = checkpoint['kmeans']

        return model

    def process_predictions(self, pred, raw_image_path, raw_image, image_transform, save_image_path):
        """After processing for prediction

        Args:
            pred (_type_): _description_
            image_transform (numpy.Array): Original image after transformation
            raw_image (Image): PIL image
            save_image_path (str): Path to save image

        Returns:
            (str): Path to the saved image
        """
        display_mode = self.opt.display_mode
        img_lvl_anom_score_tensor, pxl_lvl_anom_score = pred

        img_lvl_anom_score = float(img_lvl_anom_score_tensor.item())
        if self.opt.method == M_AMAZONE_PATCHCORE or self.opt.core_version == "1.1.0":
            pxl_lvl_score = float(torch.mean(pxl_lvl_anom_score))
        else:
            pxl_lvl_score = float((torch.mean(pxl_lvl_anom_score) + 2* torch.std(pxl_lvl_anom_score)).item())
        # define initial predict index
        predict_index = 0
        # ensemble pixel level score and image level score
        if self.opt.is_pixel_score_used == True:
            if pxl_lvl_score > self.opt.threshold_pixel_score or img_lvl_anom_score > self.opt.threshold_score:
                predict_index = 1
        else:
            if img_lvl_anom_score > self.opt.threshold_score:
                predict_index = 1

        predict_class = self.opt.class_names[predict_index]

        save_image(image_transform, pxl_lvl_anom_score, save_image_path,
                    predict_class=predict_class, img_size=raw_image.size,
                    display_mode=display_mode, pixel_score=getattr(self.opt, 'threshold_pixel_score', 0.0))

        save_txt_path = "{}.json".format(os.path.splitext(save_image_path)[0])
        with open(save_txt_path, 'w') as result_file:
            json.dump({
                "save_image_path": save_image_path,
                "score": {
                    "image_score": img_lvl_anom_score,
                    "pixel_score": pxl_lvl_score
                },
                "predict_index": predict_index,
                "class_name": predict_class,
                "bbox": [0.5, 0.5, 1, 1],
                "confidence": 0.0,
            }, result_file)

        return save_image_path
    
    @staticmethod
    def load_opt(opt_path: str):
        """load opt file and return data

        Args:
            opt_path (str): opt.json path

        Raises:
            FileNotFoundError: _description_

        Returns:
            dict: opt configuration
        """

        data = {}
        if os.path.isfile(opt_path):
            with open(opt_path, 'r') as f:
                data = EasyDict(json.load(f))
        else:
            raise FileNotFoundError(f"File {opt_path} does not exist.")

        return data
    
    @staticmethod
    def create_validation_folder(valid_folder_path):
        """
            create validation folder if it doesn't exist or not qualified
        Args:
            valid_folder_path (str): _description_

        Returns:
            Bool: True: Exist and qualified
                  False: Not exist or not qualified
        """
        # Check if the validation folder already exists
        if os.path.exists(valid_folder_path):
            # If the validation folder is not empty but contains no images, print a warning message
            files = os.listdir(valid_folder_path)
            if len(files) == 0 or all(imghdr.what(os.path.join(valid_folder_path, f)) is None for f in files):
                print("Warning: Validation folder exists but contains no image files.")
            else:
                print("Validation folder already exists.")
                return False
            
        print("========")
        print("Creating validation folder...")
        os.makedirs(valid_folder_path, exist_ok=True)
        
        return True

    @staticmethod
    def generate_validation_dataset(base_path, valid_folder_path, split_percentage):
        """
           Create a validation dataset based on training dataset
        
        Args:
            base_path (str): path to data directory
            valid_folder_path (str): validation folder path
            split_percentage (float): ratio of validation dataset / training dataset
        """
        train_path = os.path.join(base_path, "train", "good")

        # Get a list of all the files in the train folder
        file_list = os.listdir(train_path)
        num_files = int(len(file_list) * split_percentage)
        # Randomly select files to move
        files_to_copy = random.sample(file_list, num_files)
        
        # Move the selected files from train folder to the validation folder
        for file_name in files_to_copy:
            src_path = os.path.join(train_path, file_name)
            dest_path = os.path.join(valid_folder_path, file_name)
            shutil.move(src_path, dest_path)
        print("Generating validation dataset is finished!")
    
    def calculate_threshold_score(self, weight_path, valid_path):
        """
            calculate the threshold score from the validation dataset

        Args:
            weight_path (str): weight path
            valid_path (str): validation images folder path

        Returns:
            List: highest threshold score and pixel score
            List: score per image
        """
        print("========")
        print("Calculating threshold score...")

        model = self.build_model(weight_path)
        model.eval()
        transform = get_transform(SIZE)
        if self.opt.method == M_PATCHCORE_CLUSTER:
            second_transform = get_transform(SECOND_SIZE)
        all_valid_score = []
        all_valid_pixel_score = []
        result_per_img = []
        for name in os.listdir(valid_path):
            try:
                img = Image.open(valid_path + "/" + name)
                img = img.convert("RGB")
                img_trans = transform(img).unsqueeze(dim=0)

                ## predict
                if self.opt.method == M_PATCHCORE_CLUSTER:
                    sc_img_trans = second_transform(img).unsqueeze(dim=0)
                    img_score, pixel_lvl_score = model.predict(img_trans, sc_img_trans)
                else:
                    img_score, pixel_lvl_score = model.predict(img_trans)
                all_valid_score.append(img_score.numpy())
                if self.opt.method == M_AMAZONE_PATCHCORE or self.opt.core_version == "1.1.0":
                    pixel_score = torch.mean(pixel_lvl_score).numpy()
                else:
                    pixel_score = (torch.mean(pixel_lvl_score) + 2* torch.std(pixel_lvl_score)).numpy()
                all_valid_pixel_score.append(pixel_score)
                result_per_img.append({
                    "img_path": os.path.join(valid_path, name),
                    "score": {
                        "image_score": float(img_score.numpy()),
                        "pixel_score": float(pixel_score),
                    }})
            except Exception as e:
                print(e)
                continue
        
        if self.opt.method == M_AMAZONE_PATCHCORE or self.opt.core_version == "1.1.0":
            threshold = float(max(all_valid_score))
        else:
            threshold = float(np.mean(all_valid_score) + 2*np.std(all_valid_score))
        threshold_pixel_score = float(max(all_valid_pixel_score))

        return [threshold, threshold_pixel_score], result_per_img

    def save_threshold_score(self, threshold, threshold_pixel_score: float = 0.0):
        """
            save to opt.json

        Args:
            threshold (int): threshold score
            threshold_pixel_score(float): threshold for pixel level score. Default 0.0
        """
        opt_file_path = self.opt.opt_path

        if os.path.exists(opt_file_path):
            with open(opt_file_path, 'r+') as opt_file:
                opt_data = json.load(opt_file)
                opt_data['threshold_score'] = threshold
                opt_data['threshold_pixel_score'] = threshold_pixel_score
                opt_file.seek(0)
                json.dump(opt_data, opt_file)
                opt_file.truncate()
        else:
            with open(opt_file_path, 'w') as opt_file:
                json.dump({
                    'threshold_score': threshold,
                    'threshold_pixel_score': threshold_pixel_score
                }, opt_file)

        print(f"Saving threshold score = {threshold}, threshold_pixel_score = {threshold_pixel_score} to opt.json")

    def load_threshold_score(self, weight_path, valid_path):
        """
            read threshold_score file and return threshold score calculated from validation dataset if opt.json is exist
            if not exist it will calculate threshold score from validation dataset

        Args:
            weight_path (str): weight path
            valid_path (str): validation images folder path

        Returns:
            List: threshold score and pixel score
        """
        opt_file_path = self.opt.opt_path

        if os.path.exists(opt_file_path):
            with open(opt_file_path, 'r') as opt_file:
                opt_data = json.load(opt_file)

                if 'threshold_score' in opt_data:
                    if self.opt.is_pixel_score_used == True and 'threshold_pixel_score' in opt_data:

                        return [opt_data['threshold_score'], opt_data['threshold_pixel_score']]
                    else:
                        return [opt_data['threshold_score'], 0.0]

        if os.path.exists(opt_file_path) or not 'threshold_score' in opt_data:
            list_threshold, _ = self.calculate_threshold_score(weight_path, valid_path)
            self.save_threshold_score(list_threshold[0], list_threshold[1])
            return list_threshold

    @staticmethod
    def create_augment_images(good_folder_path, image_files, number_of_copy):
        """
           Augment the images in a directory randomly, then generate training dataset

        Args:
            good_folder_path (str): The path to the directory containing the images
            image_files (list): list of image files
            number_of_copy (int): maximum number of copies per image to be augmented

        """
        print("========")
        print("Generating training data...")

        if number_of_copy == 0:
            number_of_copy = 5 # set maximum number of copy to default value if dataset size is not qualified

        for file in image_files:
            file_path = os.path.join(good_folder_path, file)
            im = Image.open(file_path)
            im = im.convert("RGB")
            transform = create_transform_function()
            for index in range(number_of_copy):
                timestamp = datetime.now().strftime("%d%m%Y%H%M%S")
                trans_im = transform(im)
                image_basename = os.path.basename(file)
                new_image_path = os.path.join(
                    good_folder_path, 
                    os.path.splitext(image_basename)[0] + f"_{timestamp}" + os.path.splitext(image_basename)[1])

                trans_im.save(new_image_path)
        print("Generating training data is finished!")

    @staticmethod
    def load_dataset(good_folder_path):
        """
            load data from specified folder, return list of image files name and number of images

        Args:
            good_folder_path (str): The path to the directory containing the images

        Returns:
            number: number of images
            list: list of image files
        """

        image_files = [f for f in os.listdir(good_folder_path) if imghdr.what(os.path.join(good_folder_path, f)) is not None]
        return len(image_files), image_files

    def generate_training_data(self, train_path, number_of_copy):
        """
            Augment the images in a directory randomly, then generate training dataset

        Args:
            train_path (str): The path to the directory containing the images to be augmented
            number_of_copy (int): maximum number of copies per image to be augmented
        """

        # Find folder good to generate training data
        good_folder_path = os.path.join(train_path, "good")

        if not os.path.exists(good_folder_path):
            print("The 'good' folder does not exist.")
            return

        required_dataset_size = 10
        while True:
            dataset_size, image_files = self.load_dataset(good_folder_path)

            if dataset_size >= required_dataset_size:
                break

            print(f'Dataset has {dataset_size} images and less than required.')

            self.create_augment_images(good_folder_path, image_files, number_of_copy)

    @staticmethod
    def save_model(model, method, weight_path):
        if method == M_SPADE:
            torch.save({'model_state_dict': model.state_dict(),
                    'z_lib': model.z_lib,
                    'features_maps': model.feature_maps,
                    }, weight_path)
        elif method == M_PADIM:
            torch.save({'model_state_dict': model.state_dict(),
                        'resize': model.resize,
                        'r_indices': model.r_indices,
                        'means_reduced': model.means_reduced,
                        'E_inv': model.E_inv,
                        }, weight_path)
        elif method == M_PATCHCORE or method == M_SA_PATCHCORE:
            torch.save({'model_state_dict': model.state_dict(),
                        'resize': model.resize,
                        'patch_lib': model.patch_lib,
                        }, weight_path)
        elif method == M_AMAZONE_PATCHCORE:
            model.save_weight(weight_path)
        elif method == M_PATCHCORE_FAISS:
            torch.save({'model_state_dict': model.state_dict(),
                        'resize': model.resize,
                        'patch_lib': model.patch_lib,
                        'anomaly_scorer_index': model.anomaly_scorer.nn_method.search_index,
                        }, weight_path)
        elif method == M_PATCHCORE_CLUSTER:
            torch.save({'model_state_dict': model.state_dict(),
                        'resize': model.resize,
                        'second_resize': model.second_resize,
                        'patch_lib': model.patch_lib,
                        'kmeans': model.kmeans
                        }, weight_path)

    def cross_validation(self, train_dir, kfold=3, custom_weight=None):
        threshold_score_cross = []
        threshold_pixel_cross = []
        img_glob = glob(os.path.join(train_dir, "*"))
        kf = KFold(n_splits=kfold, shuffle=True, random_state=42)
        for train_index, valid_index in kf.split(img_glob):
            train_data = [img_glob[index] for index in train_index]
            valid_data = [img_glob[index] for index in valid_index]
            transform = get_transform(SIZE)
            train_ds = TrainDatasetListPath(train_data, SIZE, transform)
            train_dl = DataLoader(train_ds)
            if custom_weight is not None:
                model = self.build_model(custom_weight)
            else:
                model = self.generate_instance_model(self.opt.method)
            model.fit(train_dl)
            all_valid_score = []
            all_valid_pixel_score = []
            for valid_img in valid_data:
                try:
                    img = Image.open(valid_img)
                    img = img.convert("RGB")
                    img_trans = transform(img).unsqueeze(dim=0)

                    ## predict
                    img_score, pixel_lvl_score = model.predict(img_trans)
                    all_valid_score.append(img_score.numpy())
                    if self.opt.method == M_AMAZONE_PATCHCORE:
                        pixel_score = torch.mean(pixel_lvl_score).numpy()
                    else:
                        pixel_score = (torch.mean(pixel_lvl_score) + 2* torch.std(pixel_lvl_score)).numpy()
                    all_valid_pixel_score.append(pixel_score)
                except Exception as e:
                    print(e)
                    continue
            if self.opt.method == M_AMAZONE_PATCHCORE:
                threshold = float(max(all_valid_score))
            else:
                threshold = float(np.mean(all_valid_score) + 2*np.std(all_valid_score))
            threshold_pixel_score = float(max(all_valid_pixel_score))
            # append a threshold fold
            threshold_score_cross.append(threshold)
            threshold_pixel_cross.append(threshold_pixel_score)
        final_threshold_score = float(np.mean(threshold_score_cross))
        final_threshold_pixel = float(np.mean(threshold_pixel_cross))
        return final_threshold_score, final_threshold_pixel