import torch
import scipy.io as sio
from argparse import Namespace

import numpy as np
import streamlit as st
import matplotlib.pyplot as plt
from PIL import Image, ImageEnhance

import multispectral as mspec

def GetRGB( snp, scale=2 ):
    try: 
        snp = np.transpose( snp, [ 1, 2, 0 ] ).astype( float )
    except: 
        pass # not a multi-channel image
    snp = np.round( 255*( snp - snp.min() ) / ( snp.max() - snp.min() ) ).astype( np.uint8 )
    im = Image.fromarray( snp )
    im = im.resize( ( im.size[0]//scale, im.size[1]//scale ), resample=Image.Resampling.NEAREST )
    if np.ndim( snp ) == 3: 
        enhancer = ImageEnhance.Brightness( im )
        im = enhancer.enhance( 5 )
    return im

def plot_zoomed_region( image, x1, y1, x2, y2, title='Zoomed in region' ):
    # Crop the image to the zoom box coordinates
    cropped_image = image.crop( ( x1, y1, x2, y2 ) )

    # Convert cropped image to numpy array for plotting
    img_array = np.array( cropped_image )

    # Plot the cropped image
    fig, ax = plt.subplots()
    ax.imshow( img_array )
    ax.set_title( title )
    ax.axis( 'off' )
    return fig

def dashboard_header():
    cultural_reference = 'https://www.youtube.com/watch?v=ACmydtFDTGs'
    st.markdown( '# `hotdog-nothotdog v2.0`' )
    st.markdown( '## - Jian Yang (possibly)' )
    st.markdown( 'Cultural reference: [this clip](%s) from _Silicon Valley_. '%cultural_reference )
    return

def dashboard_intro():
    st.markdown( 
        '''
        # TL;DR
        An interactive, trained-from-scratch semantic segmenter for 6-channel multi-spectral imagery data, trained solely on the RIT-18 dataset. 
        This is a work in progress. 
        <a href="#app">Go to the app. </a>

        # Dataset description
        RIT-18 consists of 19 classes within the scan range of a multi-spectral drone camera over a New York state park. It is highly unbalanced.  

        # Training methodology
        This model was trained repeatedly on randomized $256 \\times 256$ pixel slices of the same swathe region of the drone, but with the pixel labels at each iteration increasingly diverse. 
        For example, for the very first iteration, the dataset was relabeled from 19 classes to only "Vegetation" and "Other" (_i.e._, a binary segmenter). 
        The resulting model became the starting point for the training of the next iteration, in which the pixels with the "Other" label were further re-labeled "Water" and "Other", for a total of 3 classes. 
        Similarly, in later iterations, "Vegetation" will be broken down to "Trees" and "Grass" and so on. 
        In this manner, the goal is to obtain a final model trained on all 19 classes of the RIT-18 dataset.

        # Goal 
        To see if it's possible to incrementally train the model in this manner. 
        If so, how influential are the near infrared (NIR) bands in the segmentation outcome?  


           - **Dataset**: [RIT-18](https://arxiv.org/abs/1703.06452)
           - **Model architecture**: U-net

        # The app<a name="app">
        ''', 
        unsafe_allow_html=True
     )
    return

def get_model( num_classes, model_state ):
    mynn = mspec.MultiSpectNet( 
        img_size=[ 256, 256 ], 
        num_channels=7, 
        num_classes=num_classes, 
        learning_rate=1.e-3
    )
    mynn.model.load_state_dict( torch.load( model_state ) )
    return mynn