import streamlit as st
import pandas as pd
import os
import time
from datetime import datetime

from alphabase.yaml_utils import save_yaml

from peptdeep.settings import global_settings
from peptdeep.webui.ui_utils import (
    get_posix, select_files, file_type_selectbox
)

from peptdeep.webui.server import queue_folder

def nce_search():
    grid_nce_first = st.number_input('Start NCE for grid NCE search',value = global_settings['model_mgr']['transfer']['grid_nce_first']*1.0,step = 1.0)
    global_settings['model_mgr']['transfer']['grid_nce_first'] = grid_nce_first
    grid_nce_last = st.number_input('End NCE for grid NCE search',min_value = grid_nce_first, value = global_settings['model_mgr']['transfer']['grid_nce_last']*1.0,step = 1.0)
    global_settings['model_mgr']['transfer']['grid_nce_last'] = grid_nce_last
    grid_nce_step = st.number_input('Step NCE for grid NCE search', value = global_settings['model_mgr']['transfer']['grid_nce_step']*1.0,step = 1.0)
    global_settings['model_mgr']['transfer']['grid_nce_step'] = grid_nce_step

    grid_instrument = st.multiselect('Instruments for grid NCE search', (
        global_settings['model_mgr']['instrument_group']
    ),default = global_settings['model_mgr']['transfer']['grid_instrument']) 
    global_settings['model_mgr']['transfer']['grid_instrument'] = grid_instrument

def fine_tune():
    epoch_ms2 = st.number_input('Epoch to train MS2 model', value = global_settings['model_mgr']['transfer']['epoch_ms2'])
    global_settings['model_mgr']['transfer']['epoch_ms2'] = epoch_ms2
    warmup_epoch_ms2 = st.number_input('Warmup epoch to train MS2 model', value = global_settings['model_mgr']['transfer']['warmup_epoch_ms2'], max_value=epoch_ms2)
    global_settings['model_mgr']['transfer']['warmup_epoch_ms2'] = warmup_epoch_ms2
    batch_size_ms2 = st.number_input('Mini-batch size to train MS2 model', value = global_settings['model_mgr']['transfer']['batch_size_ms2'])
    global_settings['model_mgr']['transfer']['batch_size_ms2'] = batch_size_ms2
    lr_ms2 = st.number_input('Learning rate to train MS2 model', value = global_settings['model_mgr']['transfer']['lr_ms2'], format='%e', step=1e-5)
    global_settings['model_mgr']['transfer']['lr_ms2'] = lr_ms2
    
    epoch_rt_ccs = st.number_input('Epoch to train RT and CCS models', value = global_settings['model_mgr']['transfer']['epoch_rt_ccs'])
    global_settings['model_mgr']['transfer']['epoch_rt_ccs'] = epoch_rt_ccs
    warmup_epoch_rt_ccs = st.number_input('Warmup epoch to train RT and CCS model', value = global_settings['model_mgr']['transfer']['warmup_epoch_rt_ccs'], max_value=epoch_rt_ccs)
    global_settings['model_mgr']['transfer']['warmup_epoch_rt_ccs'] = warmup_epoch_rt_ccs
    batch_size_rt_ccs = st.number_input('Mini-batch size to train RT and CCS model', value = global_settings['model_mgr']['transfer']['batch_size_rt_ccs'])
    global_settings['model_mgr']['transfer']['batch_size_rt_ccs'] = batch_size_rt_ccs
    lr_rt_ccs = st.number_input('Learning rate to train RT and CCS model', value = global_settings['model_mgr']['transfer']['lr_rt_ccs'], format='%e', step=1e-5)
    global_settings['model_mgr']['transfer']['lr_rt_ccs'] = lr_rt_ccs

def show():
    st.write("# Transfer learning")

    model_output_folder = st.text_input(
        'Model output folder', 
        value=global_settings['model_mgr']['transfer']['model_output_folder'].format(
            PEPTDEEP_HOME=global_settings['PEPTDEEP_HOME']
        )
    )
    model_output_folder = get_posix(
        model_output_folder
    )
    global_settings['model_mgr']['transfer']['model_output_folder'] = model_output_folder

    st.write("### PSM files for training")
    psm_type = file_type_selectbox(
        ui_label='PSM type',
        st_key='select_psm_type',
        default_type=global_settings['model_mgr']['transfer']['psm_type'],
        monitor_files=global_settings['model_mgr']['transfer']['psm_files'],
        choices=global_settings['model_mgr']['transfer']['psm_type_choices'],
        index=global_settings['model_mgr']['transfer']['psm_type_choices'].index(
            global_settings['model_mgr']['transfer']['psm_type']
        )
    )
    psm_type_to_ext_dict = {
        "alphapept": ".ms_data.hdf",
        "pfind": ".spectra",
        "maxquant": "msms.txt",
        "diann": "tsv",
        "speclib_tsv": "tsv",
    }
    global_settings['model_mgr']['transfer']['psm_type'] = psm_type
    select_files(
        global_settings['model_mgr']['transfer']['psm_files'], 
        psm_type_to_ext_dict[psm_type], 
        "Input PSM files"
    )

    st.write("### MS files for training")
    ms_file_type = file_type_selectbox(
        ui_label='MS file type',
        st_key='select_ms_file_type',
        default_type=global_settings['model_mgr']['transfer']['ms_file_type'],
        monitor_files=global_settings['model_mgr']['transfer']['ms_files'],
        choices=global_settings['model_mgr']['transfer']['ms_file_type_choices'], 
        index=global_settings['model_mgr']['transfer']['ms_file_type_choices'].index(
            global_settings['model_mgr']['transfer']['ms_file_type']
        )
    )
    ms_type_to_ext_dict = {
        "alphapept_hdf": ".ms_data.hdf",
        "thermo_raw": ".raw",
        "mgf": ".mgf",
        "mzml": ".mzml",
    }
    global_settings['model_mgr']['transfer']['ms_file_type'] = ms_file_type
    select_files(
        global_settings['model_mgr']['transfer']['ms_files'], 
        ms_type_to_ext_dict[ms_file_type], 
        "Input MS files"
    )

    st.write("### Training settings")
    fine_tune()

    psm_num_to_train_ms2 = st.number_input('PSM num to tune MS2 model', value = int(global_settings['model_mgr']['transfer']['psm_num_to_train_ms2']), step = 1)
    global_settings['model_mgr']['transfer']['psm_num_to_train_ms2'] = psm_num_to_train_ms2
    psm_num_per_mod_to_train_ms2 = st.number_input('PSM num per mod to tune MS2 model', value = int(global_settings['model_mgr']['transfer']['psm_num_per_mod_to_train_ms2']), step = 1)
    global_settings['model_mgr']['transfer']['psm_num_per_mod_to_train_ms2'] = psm_num_per_mod_to_train_ms2
    psm_num_to_train_rt_ccs = st.number_input('PSM num to tune RT and CCS model', value = int(global_settings['model_mgr']['transfer']['psm_num_to_train_rt_ccs']), step = 1)
    global_settings['model_mgr']['transfer']['psm_num_to_train_rt_ccs'] = psm_num_to_train_rt_ccs
    psm_num_per_mod_to_train_rt_ccs = st.number_input('PSM num per mod to tune RT and CCS model', value = int(global_settings['model_mgr']['transfer']['psm_num_per_mod_to_train_rt_ccs']), step = 1)
    global_settings['model_mgr']['transfer']['psm_num_per_mod_to_train_rt_ccs'] = psm_num_per_mod_to_train_rt_ccs
    top_n_mods_to_train = st.number_input('Top n mods to tune', value = int(global_settings['model_mgr']['transfer']['top_n_mods_to_train']), step = 1)
    global_settings['model_mgr']['transfer']['top_n_mods_to_train'] = top_n_mods_to_train


    st.write('### Grid search for NCEs and instruments')
    st.write('If NCE and instrument are unknown, grid search will look for the best NCE and instrument)')
    grid_nce_search = st.checkbox('Enabled', global_settings['model_mgr']['transfer']['grid_nce_search'])
    global_settings['model_mgr']['transfer']['grid_nce_search'] = grid_nce_search
    if grid_nce_search is True:
        nce_search()

    now = datetime.now()
    current_time = now.strftime("%Y-%m-%d--%H-%M-%S.%f")
    task_name = st.text_input("Task name", value=f"peptdeep_transfer_{current_time}")
    
    if st.button('Submit for transfer learning'):
        global_settings['task_type'] = 'train'

        if not os.path.exists(model_output_folder):
            os.makedirs(model_output_folder)

        yaml_path = f'{queue_folder}/{task_name}.yaml'
        save_yaml(
            yaml_path, global_settings
        )
        save_yaml(
            os.path.join(
                model_output_folder, 
                f'{task_name}.yaml'
            ), 
            global_settings
        )
        
        st.write(f'`train` task saved as "{os.path.expanduser(yaml_path)}" in the task queue')