import streamlit as st
st.set_page_config(layout="wide")
import pandas as pd
import numpy as np
from zipfile import ZipFile

import plotly.express as px
import plotly.graph_objs as go

LLR_FILE='ALL_hum_isoforms_ESM1b_LLR.zip'

with ZipFile(LLR_FILE) as myzip:
                uids=[i.split('/')[-1].split('_')[0] for i in myzip.namelist()][1:]

df=pd.read_csv('isoform_list.csv',index_col=0)

def load_LLR(uniprot_id):
  '''Loads the LLRs for a given uniprot id. Returns a 20xL dataframe 
     rows are indexed by AA change, 
     (AAorder=['K','R','H','E','D','N','Q','T','S','C','G','A','V','L','I','M','P','Y','F','W'])
     columns indexed by WT_AA+position e.g, "G 12"
     Usage example: load_LLR('P01116') or load_LLR('P01116-2')'''
  with ZipFile(LLR_FILE) as myzip:
    data = myzip.open(myzip.namelist()[0]+uniprot_id+'_LLR.csv')
  return pd.read_csv(data,index_col=0)


def plot_interactive(uniprot_id):
  primaryLLR = load_LLR(uniprot_id)
  
  template='plotly_white'

  fig = px.imshow(primaryLLR.values, x=primaryLLR.columns, y=primaryLLR.index, color_continuous_scale='Viridis_r',zmax=0,zmin=-20,
                  labels=dict(y="Amino acid change", x="Protein sequence", color="LLR"),
                  template=template,
                  title='uniprot_id: '+uniprot_id)
  fig.update_xaxes(tickangle=-90,range=[0,99],rangeslider=dict(visible=True),dtick=1)
  fig.update_yaxes(dtick=1)
  fig.update_layout({
  'plot_bgcolor': 'rgba(0, 0, 0, 0)',
  'paper_bgcolor': 'rgba(0, 0, 0, 0)',
  },font={'family':'Arial','size':11},
  hoverlabel=dict(font=dict(family='Arial', size=14)))

  fig.update_traces(
      hovertemplate="<br>".join([
          "<b>%{x} %{y}</b>"+
          " (%{z:.2f})", 
      ])+'<extra></extra>'
  )
  
  return fig

selection = st.selectbox("uniprot_id:", df)
uid=df[df.txt==selection].index.values[0]

fig = plot_interactive(uid)
fig.update_layout(width = 800, height = 600, autosize = False)
st.plotly_chart(fig, use_container_width=True)