import sys
import time
from data_setup import Dataset, TrainDatasetFolder
from model_builder import SPADE, PaDiM, PatchCore
from utils import print_and_export_results

from typing import List
from datetime import datetime

import torch
from torch.utils.data import DataLoader
import random
import numpy as np

# seeds
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)

import warnings # for some torch warnings regarding depreciation
warnings.filterwarnings("ignore")

ALLOWED_METHODS = ["spade", "padim", "patchcore"]
method = sys.argv[1]
train_dir = f"../data/ichiwa/train"
#CONSTANT
K_SPADE = 1
D_REDUCED = 350
F_CORESET = .10
BACKBONE_NAME = "wide_resnet50_2"

def run_model(method: str, train_dir: str):
    results = {}

    if method == "spade":
        model = SPADE(
            k=K_SPADE,
            backbone_name=BACKBONE_NAME,
        )
    elif method == "padim":
        model = PaDiM(
            d_reduced=D_REDUCED,
            backbone_name=BACKBONE_NAME,
        )
    elif method == "patchcore":
        model = PatchCore(
            f_coreset=F_CORESET, 
            backbone_name=BACKBONE_NAME,
        )

    print(f"\n█│ Running {method} on dataset.")
    print(  f" ╰{'─'*(len(method)+23)}\n")

    train_ds = TrainDatasetFolder(train_dir,224)
    train_dl = DataLoader(train_ds)

    print("     Training ...")
    start = time.perf_counter()

    model.fit(train_dl)

    timestamp = datetime.now().strftime("%d_%m_%Y_%H_%M_%S")
    name = f"./results/{method}_{timestamp}.pt"
    if method == "spade":
        torch.save({'model_state_dict': model.state_dict(),
                    'z_lib': model.z_lib,
                    'features_maps': model.feature_maps,
                    }, name )
    elif method == "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,
                    }, name )
    elif method == "patchcore":
        torch.save({'model_state_dict': model.state_dict(),
                    'resize': model.resize,
                    'patch_lib': model.patch_lib,
                    }, name )
             
    print("     DONE!")
    end = time.perf_counter()
    print(f"[INFO] Execution time: {end - start : .3f}s")





if __name__ == "__main__":
    method = method.lower()
    assert method in ALLOWED_METHODS, f"Select from {ALLOWED_METHODS}."
    print("========")
    run_model(method, train_dir)