import plotly.graph_objects as go
import src.utils as utils
from src.heatmap.dash_utils import *
import plotly.express as px
import os
import base64
import io
from PIL import Image
import numpy


def get_heatmap(path, keys_yaml, show_setting_id=True, only_factors=False, mode='samples', labels='factors', background=False, sample_labels=False, condition_labels=False, transpose=False, drop_defaults=False):
    settings, experimental_factors, organisms, max_vals, options_pretty, annotated_dict, max_annotation, conditions = get_data(path, keys_yaml, mode=mode, drop_defaults=drop_defaults, transpose=transpose)

    techniques = list(utils.find_keys(path, 'techniques'))
    setting_techniques = {}
    if len(techniques) > 0:
        techniques = techniques[0]
        for elem in techniques:
            try:
                sett_id = elem['setting']
                try:
                    technique = elem['technique']
                    setting_techniques[sett_id] = technique
                except KeyError:
                    pass
            except KeyError:
                pass

    heatmaps=[]

    for value in settings:
        fig = None
        if len(settings[value]) > 0:

            options = [key.replace('_num', '') for key in settings[value] if key.endswith('_num')]
            
            if only_factors:
                sorter = experimental_factors[value]
            else:
                sorter = experimental_factors[value] + [o for o in options if o not in experimental_factors[value]]
            df_empty = settings[value].dropna(axis=1, how='all')
            sorter = [x for x in sorter if x in df_empty.columns]
            val_sorter = [0]
            for elem in sorter:
                if elem in max_annotation[value]:
                    val_sorter.append(max(val_sorter[-1]+(50 if not transpose else 150), val_sorter[-1]+(15*max_annotation[value][elem])+20))
                else:
                    val_sorter.append(val_sorter[-1]+(50 if not transpose else 150))

            df = [settings[value][f'{key}_num'] for key in sorter]   
            annotated = [annotated_dict[value][key] for key in sorter if key in annotated_dict[value]]
            if transpose:
                annotated = numpy.transpose(annotated)
            option_text = []
            label_text = []
            for key in sorter:
                if labels == 'all' or (labels == 'factors' and key in annotated_dict[value] and key in experimental_factors[value]):
                    label_vals = []
                    for elem in annotated_dict[value][key]:
                        splitted = elem.split('<br>')
                        new_val = []
                        for val in splitted:
                            if len(val) > 20:
                                new_val.append(f'{val[:18]}...')
                            else:
                                new_val.append(val)
                        label_vals.append('<br>'.join(new_val))
                    label_text.append(label_vals)
                else:
                    label_text.append(['']* len(annotated_dict[value][key]))

            for option in sorter:
                if option in experimental_factors[value]:
                    option_text.append(color("red", f'<b>{options_pretty[option]}</b>'))
                else:
                    option_text.append(color("black", options_pretty[option]))
                    
            plotly_colors = []

            colors_scales = [px.colors.qualitative.G10, px.colors.qualitative.Plotly, px.colors.qualitative.D3, px.colors.qualitative.T10]
            
            i = 0
            while len(plotly_colors) < max_vals[value]:
                plotly_colors += [x for x in colors_scales[i%len(colors_scales)] if x not in ["#DC3912", "#EF553B", "#D62728", "#F58518"]]
                i += 1

            colors = [[0, 'white']]
            for i in range(1, max_vals[value]+1):
                colors.append([i*1/(max_vals[value]), plotly_colors[i-1]])
                if i*1/(max_vals[value]) != 1:
                    colors.append([((i*1)/(max_vals[value]))+0.001, 'white'])

            if transpose:
                if mode == 'samples':
                    if condition_labels:
                        if sample_labels:
                            
                            y_value = [settings[value]['condition_label'], settings[value]['sample_label']]
                        else:
                            y_value = [settings[value]['condition_label'], settings[value]['sample_index']]
                    else:
                        if sample_labels:
                            y_value = [settings[value]['condition_index'], settings[value]['sample_label']]
                        else:
                            y_value = [settings[value]['condition_index'], settings[value]['sample_index']]
                else:
                    if condition_labels:
                        y_value = settings[value]['condition_label']
                    else:
                        y_value = settings[value]['condition_index']
                        
                x_value = val_sorter
                label_text = numpy.transpose(label_text)
                df = numpy.transpose(df)
            else:
                if mode == 'samples':
                    if condition_labels:
                        if sample_labels:
                            x_value = [settings[value]['condition_label'], settings[value]['sample_label']]
                        else:
                            x_value = [settings[value]['condition_label'], settings[value]['sample_index']]
                    else:
                        if sample_labels:
                            x_value = [settings[value]['condition_index'], settings[value]['sample_label']]
                        else:
                            x_value = [settings[value]['condition_index'], settings[value]['sample_index']]
                else:
                    if condition_labels:
                         x_value = settings[value]['condition_label']
                    else:
                        x_value = settings[value]['condition_index']
                y_value = val_sorter
                                
            heatmap = [go.Heatmap(
                        z=df,
                        zmin=0,
                        zmax=max_vals[value],
                        x=x_value,
                        y=y_value,
                        showscale=False,
                        customdata=annotated,
                        text=label_text,
                        texttemplate="%{text}",
                        hovertemplate = "%{customdata}",
                        hoverongaps = False,
                        colorscale = colors,
                        ),
                    ]
            
            data_input = heatmap
            
            if mode == 'samples':
                if sample_labels:
                    sample_label_height = max(settings[value]['sample_label_height'])
                else:
                    sample_label_height = 1
            else: sample_label_height = 0

            if condition_labels:
                condition_label_height = max(settings[value]['condition_label_height'])
            else:
                condition_label_height = 1

            my_cell_width = 150
            #my_cell_height = max(50, (15*max_annotation[value])+20)
            if transpose:
                top_margin = 100
                bottom_margin = 20
                left_margin = (150 if condition_labels else 100) + (150 if sample_labels else 100)
                right_margin = 20
                my_width = val_sorter[-1]
                my_height = 50 * (len(settings[value]["sample_index"]) if mode=='samples' else len(settings[value]["condition_index"]))
            else:
                top_margin = 15*sample_label_height + 15*condition_label_height + 70 # minimum 100
                bottom_margin = 20
                left_margin = 200
                right_margin = 20
                my_height = val_sorter[-1] #my_cell_height*len(sorter)

                if mode=='samples':
                    my_width = my_cell_width*len(settings[value]["sample_index"])
                else:
                    my_width = my_cell_width*len(settings[value]['condition_index'])
            organism_path = os.path.join(os.path.dirname(__file__), '..', '..', 'images', f'{organisms[value]}.png')
            images = None

            if os.path.isfile(organism_path):
                plotly_logo = base64.b64encode(open(organism_path, 'rb').read())
                imgdata = base64.b64decode(plotly_logo)
                im = Image.open(io.BytesIO(imgdata))
                im_width, im_height = im.size
                y_side = 100
                my_ysize=y_side/my_height
                x_side = y_side*im_width / im_height
                my_xsize = x_side/my_width

                left_margin = max(left_margin, x_side)

                images = [dict(
                    source='data:image/png;base64,{}'.format(plotly_logo.decode()),
                    xref="paper", yref="paper",
                    x=0, y=1,
                    sizey=my_ysize, sizex=my_xsize,
                    xanchor="right", yanchor="bottom"
            )]

            between = []
            technique_list = setting_techniques[value] if value in setting_techniques else None
            if technique_list is not None:
                technique_list = f' - {", ".join(technique_list)}'
            else:
                technique_list = ''
            for i in range(len(val_sorter)-1):
                between.append((val_sorter[i+1]+val_sorter[i])/2)
            layout = go.Layout(
                images = images,
                height=my_height + top_margin + bottom_margin,
                width=my_width + left_margin + right_margin,
                margin=dict(l=left_margin, r=right_margin, t=top_margin, b=bottom_margin),
                autosize=False,
                title=dict(
                    text=f'<b>Setting {value}{technique_list}</b>' if show_setting_id else '', 
                    font=dict(size=15), 
                    automargin=False,
                    yref='container',
                    x=(left_margin+(0.1*150))/(my_width+left_margin+right_margin),
                    y=1-(30/(my_height+top_margin+bottom_margin)*100)/(my_height+top_margin+bottom_margin),
                    xanchor="left", yanchor="top",
                    subtitle=dict(
                        text=f'Organism: {organisms[value]}',
                        font=dict(size=14, lineposition='under'),
                        ),
                    ),
                plot_bgcolor='rgb(255, 255, 255)' if background else 'rgba(0,0,0,0)',
                paper_bgcolor='rgb(255, 255, 255)' if background else 'rgba(0,0,0,0)',
                xaxis=dict(
                    side="top",
                    automargin=False,
                ) if not transpose else dict(
                    tickmode="array",
                    ticktext=option_text, 
                    tickvals=between,
                    side="top",
                    automargin=False
                ),
                yaxis=dict(
                    tickmode="array",
                    ticktext=option_text, 
                    tickvals=between,
                    autorange="reversed",
                    automargin=False
                ) if not transpose else dict(
                    automargin=False,
                    autorange="reversed",
                )
            )    

            fig = go.Figure(data=data_input, layout=layout)     
            if transpose:
                for i in val_sorter:
                    fig.add_vline(i, line_width=0.5)
                
                for i in between:
                    fig.add_vline(i, line_dash="dash", line_color='lightgrey',layer='below') 

                if mode == 'samples':
                    for i in range(len(settings[value]['sample_index'])+1):
                        fig.add_hline(i-0.5, line_width=0.5)
                else:
                    for i in range(len(settings[value]['condition_index'])+1):
                        fig.add_hline(i-0.5, line_width=0.5)
                
                fig.add_vrect(x0=0, x1=val_sorter[len(experimental_factors[value])], line=dict(color="red", width=5), layer='above')

            else:
                for i in val_sorter:
                    fig.add_hline(i, line_width=0.5)
                
                for i in between:
                    fig.add_hline(i, line_dash="dash", line_color='lightgrey',layer='below')      

                if mode == 'samples':
                    for i in range(len(settings[value]['sample_index'])+1):
                        fig.add_vline(i-0.5, line_width=0.5)
                else:
                    for i in range(len(settings[value]['condition_index'])+1):
                        fig.add_vline(i-0.5, line_width=0.5)
                
                fig.add_hrect(y0=0, y1=val_sorter[len(experimental_factors[value])], line=dict(color="red", width=5), layer='above')
            
            fig.update_layout(modebar={'remove':['zoom', 'pan', 'zoomIn', 'zoomOut', 'autoScale']})
            fig.layout.yaxis.fixedrange = True

        heatmaps.append((value, fig, conditions[value] if len(list(conditions[value].keys())) > 0 else None))
    return heatmaps


def color(color, text):
    return f"<span style='color:{str(color)}'> {str(text)} </span>"

