import streamlit as st
import pandas as pd
import requests
from dotenv import load_dotenv

from transformers import pipeline
from PIL import Image

from info import akiec, bcc, bkl, df, mel, nv, vasc, vit_base_patch_16

load_dotenv()

URL = 'https://i.stack.imgur.com/gPR77.jpg'


def download_image():
    if st.session_state.img_url:
        st.session_state['image'] = Image.open(
            requests.get(st.session_state.img_url, stream=True).raw)
    else:
        del st.session_state['image']


def file_upload():
    if st.session_state.file_upload:
        st.session_state['image'] = Image.open(st.session_state.file_upload)
    else:
        del st.session_state['image']


def cam_upload():
    if st.session_state.camera:
        st.session_state['image'] = Image.open(st.session_state.camera)
    else:
        del st.session_state['image']


if 'image' not in st.session_state:
    st.session_state['image'] = Image.open(requests.get(URL, stream=True).raw)

st.header("Skin Cancer Classifier")

with st.sidebar:
    img_upload_tab, cam_upload_tab, url_upload_tab = st.tabs(
        ['📂 Upload', '📸 CAMERA', '🔗 URL'])

    with img_upload_tab:
        uploaded_img = st.file_uploader(
            label="Upload a Skin Lesion", on_change=file_upload, key='file_upload'
        )
    with cam_upload_tab:
        camera_img = st.camera_input(
            label='Take a picture of a Skin Lesion', on_change=cam_upload, key='camera'
        )
    with url_upload_tab:
        img_url = st.text_input(
            label="Enter the Skin Lesion URL", value=URL, on_change=download_image, key="img_url"
        )

    if 'image' in st.session_state:
        st.image(st.session_state['image'])

    analyze_btn = st.button(label='Analyze Skin Lesion', type='primary', use_container_width=True, key='analyze_btn')

if 'image' in st.session_state and st.session_state['analyze_btn']:
    with st.spinner("Analyzing..."):
        pipe = pipeline("image-classification", model="sharren/vit-beta2-0.99")
        response = pipe(st.session_state['image'])

        df = pd.DataFrame(response)
        result = df.nlargest(n=1, columns='score')
        result_body = f'Model predicts: {result["label"].item()} with {result["score"].item()*100:.2f}% confidence'

        with st.expander(label=result_body, expanded=True):
            st.subheader(f':red[{result["label"].item()}] Detected')
            st.bar_chart(data=df, x='label', y='score')

        with st.expander(label="Skin lesion analyzed"):
            st.image(st.session_state['image'])

        with st.expander(label="Model Details"):
            st.markdown(body=vit_base_patch_16)

else:
    tab_1, tab_2, tab_3, tab_4, tab_5, tab_6, tab_7 = st.tabs([
        'Actinic Keratoses', 'Basal Cell Carcinoma', 'Benign Keratosis', 'Dermatofibroma', 
        'Melanoma', 'Melanocytic Nevi', 'Vascular Lesion'
    ])
    with tab_1:
        st.subheader('Actinic Keratoses')
        st.markdown(body=akiec)
    with tab_2:
        st.subheader('Basal Cell Carcinoma')
        st.markdown(body=bcc)
    with tab_3:
        st.subheader('Benign Lesions of the Keratosis Type')
        st.markdown(body=bkl)
    with tab_4:
        st.subheader('Dermatofibroma')
        st.markdown(body=df)
    with tab_5:
        st.subheader('Melanoma')
        st.markdown(body=mel)
    with tab_6:
        st.subheader('Melanocytic Nevi')
        st.markdown(body=nv)
    with tab_7:
        st.subheader('Vascular Lesion')
        st.markdown(body=vasc)