import sys
import os
from tqdm import tqdm
from datetime import datetime
from PIL import Image
import torch
from torch import tensor
from torchvision import transforms


from model_builder import SPADE, PaDiM, PatchCore
from utils.util_indad import GaussianBlur, get_coreset_idx_randomp, get_tqdm_params, save_image
from data_transform import transform_default

# CONSTANT
K_SPADE = 1
D_REDUCED = 350
F_CORESET = .10
BACKBONE_NAME = "wide_resnet50_2"

## input weight
weight_path = sys.argv[2]

## input folder
predict_path = sys.argv[3]

## input method
ALLOWED_METHODS = ["spade", "padim", "patchcore"]
method = sys.argv[1]
method = method.lower()
assert method in ALLOWED_METHODS, f"Select from {ALLOWED_METHODS}."
if method == "spade":
    model = SPADE(
        k=K_SPADE,
        backbone_name=BACKBONE_NAME,
        )
    checkpoint = torch.load(weight_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.z_lib = checkpoint['z_lib']
    model.feature_maps = checkpoint['features_maps']
elif method == "padim":
    model = PaDiM(
        d_reduced=D_REDUCED,
        backbone_name=BACKBONE_NAME,
    )
    checkpoint = torch.load(weight_path)
    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 == "patchcore":
    model = PatchCore(
        f_coreset=F_CORESET, 
        backbone_name=BACKBONE_NAME,
    )
    checkpoint = torch.load(weight_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.resize = checkpoint['resize']
    model.patch_lib = checkpoint['patch_lib']


def predict_img(img_path:str, model, f_output):
    transform = transform_default
    img = Image.open(img_path)
    img_trans = transform(img).unsqueeze(dim=0)
    ## predict
    img_lvl_anom_score, pxl_lvl_anom_score = model.predict(img_trans)
    
    save_image(img_trans,pxl_lvl_anom_score,f_output)
    print(img_lvl_anom_score)

if __name__ == "__main__":
    ## predict from image path
    # predict_path = "../data/ichiwa/test/over_flow/20221022_160953_0.jpg"
    # predict_img(predict_path,model)

    ## predict from image folder

    timestamp = datetime.now().strftime("%d_%m_%Y_%H_%M_%S")
    f_output= f"./results/{timestamp}"
    os.mkdir(f_output)
    for name in os.listdir(predict_path):
        predict_img(predict_path+"/"+name, model, f_output)


