import streamlit as st
import PIL
import sys
import zipfile
from fastai import *
import numpy as np
import pandas as pd 
import os
import timm
from timm import create_model 
from fastai.vision.all import *
import torchvision.transforms as T

torch.device('cpu')
def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
seed_everything(42)

test = {
    'Id': ['2022-12-21 00.49.46.jpg', '2022-12-21 00.49.46.jpg'],
    'Eyes': [0.3, 0.2],
}
train_df = pd.DataFrame(test)
print(train_df)
dls = DataBlock(blocks=(ImageBlock, CategoryBlock),
                get_x=ColReader('Id'),
                get_y=ColReader('Eyes'),
                splitter=RandomSplitter(0.2),
                item_tfms=Resize(224),
                batch_tfms=setup_aug_tfms([Brightness(), Contrast(), Hue(), Saturation(), Flip(size=224)]),
               )

paw_dls = dls.dataloaders(train_df, batch_size=8, seed=12, device='cpu')
test = paw_dls.test_dl(train_df)
learn = cnn_learner(paw_dls, models.resnet50, pretrained=False, metrics=error_rate) 
learn.to('cpu')
catanddog = learn.load('dog')

st.title("PawPularity Score")

file_name = st.file_uploader("Upload a photo of your pet a get an estimate how popular it is gonna be!", type=["jpg", "jpeg"])
my_bar = st.progress(0, text="PawScore")

if file_name is not None:
    col1, col2 = st.columns(2)


    image = Image.open(file_name)
    print(image)
    img_fastai = np.array(image)
    col1.image(image, use_column_width=True)
    pred = catanddog.predict(img_fastai)[1]
    print(pred)
    def metric_rmse(input,target):
        return 100*torch.sqrt(F.mse_loss(F.sigmoid(input.flatten()), target))
    model = create_model('swin_large_patch4_window7_224', pretrained=False, num_classes=1)
    learn = Learner(paw_dls, model, loss_func = BCEWithLogitsLossFlat(),  metrics=metric_rmse)
    # learner = learner.load(f'/kaggle/input/puredog/pure-dog-{fold}')
    if pred == 0:
        learn.load('pure-dog-5')
        pred = learn.predict(image)
        col2.header("Score for this doggo:")
        col2.subheader(f"{ round(float(pred[2]) * 100, 1)}%")
        my_bar.progress(round(float(pred[2]) * 100), text="PawScore")

    if pred == 1:
        learn.load(f'pure-cat-0')
        pred = learn.predict(image)
        print(pred)
        print("Score for this catto:", min(int(pred[2] * 100 + 35), 100))
        col2.header("Score for this catto:")
        col2.subheader(f"{ round(float(pred[2]) * 100, 1)}%")
        my_bar.progress(round(float(pred[2]) * 100), text="PawScore")