import cv2
import torch
import numpy as np

from trainer import (
    extract_features,
    Encoder,
    Decoder,
    DEVICE,
    INPUT_SIZE
)



encoder=Encoder().to(DEVICE)
decoder=Decoder().to(DEVICE)


encoder.load_state_dict(
    torch.load("encoder.pt")
)

decoder.load_state_dict(
    torch.load("decoder.pt")
)


encoder.eval()
decoder.eval()



def generate_grade(
        image_path,
        warmth=0):


    f=extract_features(
        image_path
    )


    x=torch.tensor(
        f
    ).float().to(DEVICE)


    with torch.no_grad():

        z=encoder(
            x.unsqueeze(0)
        )

        result=decoder(z)


    grade=result.cpu().numpy()[0]


    return grade



def apply_grade(
        image_path,
        grade):


    img=cv2.imread(
        image_path
    )

    hsv=cv2.cvtColor(
        img,
        cv2.COLOR_BGR2HSV
    )


    hsv[:,:,1]=np.clip(
        hsv[:,:,1]
        *
        (1+grade.mean()),
        0,
        255
    )


    return cv2.cvtColor(
        hsv,
        cv2.COLOR_HSV2BGR
    )