"""
File Name: ploter.py
Author: Sylvain Meignier
Organization: Le Mans Université, LIUM (https://lium.univ-lemans.fr/)
Creation Date: 2023-12-03
Description: plotter panel
This file is governed by the terms of the GNU General Public License v3.0.
Please see the LICENSE file for more details.
"""
import numpy as np
import pandas as pd

import param
import panel as pn
from panel.viewable import Viewer

from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, FixedTicker, Label
from bokeh.models import Range1d
from bokeh.palettes import Category20, Category10
from bokeh.core.enums import LegendLocation

from pyDendron.dataname import *
from pyDendron.app_logger import logger

class Ploter(Viewer):
    ANATOMY = 'Anatomy'
    ANATOMY_COLOR = {HEARTWOOD: 'black', 
                     PITH: 'blue', 
                     SAPWOOD:'red', 
                     CAMBIUM:'red',
                     BARK: 'red'}


    class ParamRePlot(param.Parameterized):
        data_type = param.Selector(default='Raw', objects=['Raw', 'Log', 'Detrend'], doc='where is the data')
        x_offset_mode = param.Selector(default='None', objects=['None', DATE_BEGIN, OFFSET], doc='Delay of the rings')
        y_offset_mode = param.Selector(default='Stack', objects=['Zero', 'Stack'], doc='y curve position')
        y_height = param.Selector(default='[0, max]', objects=['[0, max]', '[min, max]'], doc='height of a serie')
        y_delta = param.Number(default=0, doc='marge between two series')
        draw_type = param.Selector(default='Line', objects=['Line', 'Step', 'Spaghetti'], doc='') 
        cambium_estimation = param.Boolean(False, doc='Draw cambium estimation')
        show_dates = param.Boolean(False, doc='Add dates at the start and end of the serie')

    param_replot = ParamRePlot()
    
    figure_title = param.String(' ')
    figure_title_font_size = param.Integer(default=14, bounds=(1, 20), step=1)
    width = param.Integer(default=1000, bounds=(50, 4000), step=10)
    height = param.Integer(default=500, bounds=(50, 2000), step=10)
    line_width_tree = param.Number(default=0.5, bounds=(0.25, 4.0), step=0.25)
    line_width_chronology = param.Number(default=1, bounds=(0.25, 4.0), step=0.25)
    circle_radius = param.Number(default=0.5, bounds=(0.1, 5), step=0.1)
    color = param.Selector(default=ANATOMY, objects=['None', KEYCODE, ANATOMY], 
            doc=f'None: all black, {KEYCODE}: one color per {KEYCODE}, {ANATOMY}: color pith, sapwood... ')
    legend = param.Selector(default='Y axe', objects=['None', 'Y axe', 'In figure'], 
            doc=f'position of the legend')
    legend_font_size = param.Integer(default=14, bounds=(1, 20), step=1)
    legend_location = param.Selector(default='top_left', objects=LegendLocation, doc='Legend location')

    x_range_step = param.Integer(default=25, bounds=(5, 200), step=5)
    axis_marge = param.Integer(default=5, bounds=(0, 100), step=5)
    axis_font_size = param.Integer(default=10, bounds=(1, 20), step=1)
    border = param.Boolean(False, doc='Log transform')
    
    def __init__(self, ploter_name='ploter', **params):
        super(Ploter, self).__init__(**params)   
        self.ploter_name = ploter_name
        self.draw_data = None
        self.data = None
        self.x_range, self.y_range = None, None
        self.figure_pane = pn.pane.Bokeh(height=self.height, width=self.width)
        self._layout = self.figure_pane
 
    def __panel__(self):
        return self._layout
                
    @param.depends("width", watch=True)
    def _update_width(self):
        self.figure_pane.width = self.width

    @param.depends("height", watch=True)
    def _update_height(self):
        self.figure_pane.height = self.height
    
    def get_pith_optimun(self, data_len):
        return int(data_len*0.1)
 
    def prepare_data(self, data):
        #print('prepare_data')

        self.data = data

        def init_ColumnDataSource():
            return ColumnDataSource()

        def get_x_offset(row):
            if self.param_replot.x_offset_mode == 'None':
                return 0
            elif self.param_replot.x_offset_mode == DATE_BEGIN:
                return row[DATE_BEGIN]
            return row[OFFSET]

        def get_y_offset(row, cum_y_offset):
            data = row[DATA_VALUES]
            v = self.param_replot.y_delta
            if self.param_replot.draw_type == 'Spaghetti':
                v += 50
            else:
                v += np.nanmax(data) if self.param_replot.y_offset_mode == 'Stack' else 0
            cum_y_offset += v
            return v, cum_y_offset
        
        def get_values(row, info):
            values = row[DATA_VALUES]

            sapwood_offset = row[SAPWOOD]
            info[SAPWOOD] = init_ColumnDataSource()
            info[HEARTWOOD] = init_ColumnDataSource()
            
            if pd.isna(sapwood_offset) or sapwood_offset < 0:
                sapwood_offset = len(values) - 1
            info[HEARTWOOD].data['x'] = np.arange(0, sapwood_offset + 1) + info['x_offset']
            info[HEARTWOOD].data['w'] = values[:sapwood_offset + 1]
            info[HEARTWOOD].data['y'] = info[HEARTWOOD].data['w'] + info['y_offset'] 
            
            info['is_sapwood'] = not(pd.isna(sapwood_offset) or sapwood_offset < 0)
            info[SAPWOOD].data['x'] = np.arange(sapwood_offset, len(values)) + info['x_offset']
            info[SAPWOOD].data['w'] = values[sapwood_offset:]
            info[SAPWOOD].data['y'] = info[SAPWOOD].data['w'] + info['y_offset']
            
        def get_missing_ring(row, info):
            values = row[DATA_VALUES]
            begin = np.where(~np.isnan(values))[0][0]
            end = np.where(~np.isnan(values))[0][-1]
            
            #print('--> missing ring:', begin, end, len(values), values[begin], values[end])
            
            if begin > 0:
                info[MISSING_RING_BEGIN] = init_ColumnDataSource()
                info[MISSING_RING_BEGIN].data['x'] = np.arange(0, begin) + info['x_offset']
                info[MISSING_RING_BEGIN].data['w'] = [values[begin]] * begin
                info[MISSING_RING_BEGIN].data['y'] = [values[begin] + info['y_offset']] * begin 
            if end < len(values):
                info[MISSING_RING_END] = init_ColumnDataSource()
                info[MISSING_RING_END].data['x'] = np.arange(end, len(values)) + info['x_offset']
                info[MISSING_RING_END].data['w'] = [values[end]] * (len(values) - end)
                info[MISSING_RING_END].data['y'] = [values[end] + info['y_offset']] * (len(values) - end) 
            
        def get_pith(row, info):
            values = row[DATA_VALUES] 
            x_min = info['x_offset']
            i = np.where(~np.isnan(values))[0][0]
            w = values[i]
            info[PITH] = init_ColumnDataSource()
            if pd.notna(row[PITH]) and row[PITH]:
                info[PITH].data['x'] = [info['x_offset']]
                info[PITH].data['w'] = [w]
                info[PITH].data['y'] = [w + info['y_offset']]
            return x_min
                
        def get_cambium(row, info):
            info['is_cambium_estimated'] = False
            values = row[DATA_VALUES]
            x = len(values) - 1
            x_max = x + info['x_offset']

            w = values[np.where(~np.isnan(values))[0][-1]]
            info[CAMBIUM] = init_ColumnDataSource()
            info[CAMBIUM_ESTIMATED] = init_ColumnDataSource()
            info[CAMBIUM_BOUNDARIES] = init_ColumnDataSource()
            if pd.notna(row[CAMBIUM]) and row[CAMBIUM]:
                info[CAMBIUM].data['x'] = [x + info['x_offset']]
                info[CAMBIUM].data['w'] = ['NA']
                info[CAMBIUM].data['y'] = [w + info['y_offset']]
            else:
                if (CAMBIUM_ESTIMATED in row) and pd.notna(row[CAMBIUM_ESTIMATED]):
                    info['is_cambium_estimated'] = True
                    lower, estimated, upper = row[CAMBIUM_LOWER], row[CAMBIUM_ESTIMATED], row[CAMBIUM_UPPER]
                    wo = w + info['y_offset']
                    xe = estimated + info['x_offset']
                    xl = lower + info['x_offset']
                    xu = upper + info['x_offset']
                    
                    h = 5 if self.param_replot.draw_type == 'Spaghetti' else np.nanmean(values) * 0.2
                    info[CAMBIUM_BOUNDARIES].data['x'] = np.arange(xl, xu+1) + info['x_offset']
                    info[CAMBIUM_BOUNDARIES].data['w'] = np.array(['NA']*(xu-xl+1))
                    info[CAMBIUM_BOUNDARIES].data['y'] = np.array([w]*(xu-xl+1)) + info['y_offset']
                    info[CAMBIUM_ESTIMATED].data['x0'] = info[CAMBIUM_ESTIMATED].data['x1'] = [xu, xe, xl]
                    info[CAMBIUM_ESTIMATED].data['w'] = ['NA'] * 3
                    info[CAMBIUM_ESTIMATED].data['y0'] = [wo - h] * 3
                    info[CAMBIUM_ESTIMATED].data['y1'] = [wo + h] * 3
                    x_max = xu 
            return x_max

        def get_bark(row, info):
            values = row[DATA_VALUES]
            x = len(values)
            w = values[np.where(~np.isnan(values))[-1]]
            info[BARK] = init_ColumnDataSource()
            if pd.notna(row[BARK]) and row[BARK]:
                info[BARK].data['x'] = [x + info['x_offset']]
                info[BARK].data['w'] = [w]
                info[BARK].data['y'] = [w + info['y_offset']]
                
        def get_border(row, info):
            info['border'] = init_ColumnDataSource()
            info['border'].data['top'] = [info['y_max']]
            info['border'].data['bottom'] = [info['y_min']]
            info['border'].data['left'] = [info['x_min']]
            info['border'].data['right'] = [info['x_max']]
            
        def get_min_max(row, info):
            info['x_min'] = get_pith(row, info)
            info['x_max'] = get_cambium(row, info)
            info['w_min'] = np.nanmin(row[DATA_VALUES])
            info['w_max'] = np.nanmax(row[DATA_VALUES])
            info['w_mean'] = np.nanmean(row[DATA_VALUES])
            info['y_min'] = info['y_offset'] #info['w_min'] + info['y_offset']
            info['y_max'] = next_cum_y_offset
            info['y_mean'] = info['w_mean'] + info['y_offset']
            info['y_label'] = round(info['y_mean'], 3)
            
        def get_date(row, info):
            info[DATE_BEGIN] = str(row[DATE_BEGIN]) if pd.notna(row[DATE_BEGIN]) else ''
            info[DATE_END] = str(row[DATE_END]) if pd.notna(row[DATE_BEGIN]) else 'not dated'
                    
        draw = {}
        data = data.loc[data[CATEGORY].isin([CHRONOLOGY, TREE]),:]
        if self.param_replot.x_offset_mode != 'None':
            if data[self.param_replot.x_offset_mode].isna().any():
                logger.error(f"NA value(s) in {self.param_replot.x_offset_mode} column, can't draw")
                self.draw_data = None
                return
                    
        cum_y_offset = 0
        for i, (_, row) in enumerate(data.iterrows()):   
            info = {}
            info[CATEGORY] = row[CATEGORY]
            row[RAW] = row[DATA_VALUES]

            if self.param_replot.y_height == '[min, max]':
                row[DATA_VALUES] -= np.nanmin(row[DATA_VALUES])                
            if self.param_replot.draw_type == 'Spaghetti':
                row[DATA_VALUES] = np.array([100] * row[DATA_LENGTH])
            info['i'] = i
            info[KEYCODE] = row[KEYCODE]
            info['x_offset'] = get_x_offset(row)
            info[DATA_LENGTH] = row[DATA_LENGTH]
            h, next_cum_y_offset = get_y_offset(row, cum_y_offset)
            print(i, row[KEYCODE], cum_y_offset, h, next_cum_y_offset)
            info['y_offset'] = cum_y_offset
            get_values(row, info)
            get_bark(row, info)
            get_missing_ring(row, info)
            get_min_max(row, info)
            get_border(row, info)
            get_date(row, info)
            
            draw[info[KEYCODE]] = info
            cum_y_offset = next_cum_y_offset
        self.draw_data = draw
        #print('prepare_data end')


    def clear(self):
        self.figure_pane.object = None

    @param.depends('param_replot.y_delta', 'param_replot.y_height', 'param_replot.x_offset_mode', 
                   'param_replot.y_offset_mode', 'param_replot.draw_type', 
                   'param_replot.cambium_estimation', 'param_replot.show_dates', watch=True)
    def prepare_and_plot_inner(self):
        self.prepare_and_plot(data=self.data)
    
    def prepare_and_plot(self, data=None):
        try:
            self._layout.loading = True
            if (data is None) or (len(data) == 0):
                self.clear()
                self._layout.loading = False
                return
            self.prepare_data(data) 
        except Exception as inst:
            logger.error(f'plot : {inst}', exc_info=True)
        finally:
            self._layout.loading = False
        
        self.plot()
        #print('end plot')
    
    def on_x_range_step(self):
        if (self.figure_pane.object is not None) and pd.notna(self.figure_pane.object.x_range.start):
            x_min = self.figure_pane.object.x_range.start + self.axis_marge
            x_max = self.figure_pane.object.x_range.end - self.axis_marge
            self.figure_pane.object.xaxis[0].ticker = FixedTicker(ticks= np.arange(int(x_min), int(x_max), self.x_range_step))
            label = self.param_replot.x_offset_mode if self.param_replot.x_offset_mode != 'None' else f'{OFFSET}'
            self.figure_pane.object.xaxis[0].axis_label = label
        
    @param.depends('figure_title', 'figure_title_font_size', watch=True)
    def on_figure_title(self):
        if self.figure_pane.object is not None:
            self.figure_pane.object.title.text = self.figure_title
            self.figure_pane.object.title.text_font_size = str(self.figure_title_font_size) + 'px'
    
    @param.depends('axis_font_size', watch=True)
    def on_axis_font_size(self):
        if self.figure_pane.object is not None:
            self.figure_pane.object.yaxis.major_label_text_font_size = f'{self.axis_font_size}px'
            self.figure_pane.object.xaxis.major_label_text_font_size = f'{self.axis_font_size}px'
    
    @param.depends('legend_font_size', watch=True)
    def on_legend_font_size(self):
        if self.figure_pane.object is not None:
            self.figure_pane.object.legend.label_text_font_size = f'{self.legend_font_size}px'
            #print(self.figure_pane.object.legend.label_text_font_size)

    @param.depends('legend_location', watch=True)
    def on_legend_location(self):
        if self.figure_pane.object is not None:
            #print(self.legend_location, str(self.legend_location))
            self.figure_pane.object.legend.location = self.legend_location

    def on_legend(self):
        if self.figure_pane.object is not None:            
            #self.figure_pane.object.legend.visible = False
            y_labels = {}
            self.figure_pane.object.yaxis.visible = False
            for i, (keycode, info) in enumerate(self.draw_data.items()):
                y_labels[info['y_label']] = keycode if self.legend == 'Y axe' else str(i)
            self.figure_pane.object.yaxis.ticker = list(y_labels.keys())
            self.figure_pane.object.yaxis.major_label_overrides = y_labels
            self.figure_pane.object.ygrid.grid_line_color = None
            #print('on_legend, y_labels:',  y_labels)
            #print('on_legend, major_label_overrides:',  self.figure_pane.object.yaxis.major_label_overrides)
            
            if self.legend == 'In figure':
                #print(self.legend_location, str(self.legend_location))
                self.figure_pane.object.legend.location = self.legend_location #"top_left"
                self.figure_pane.object.legend.click_policy="mute"
                self.figure_pane.object.legend.visible = True
                
    def get_color(self, kind, rank):
        if self.color == self.ANATOMY:
            return self.ANATOMY_COLOR[kind]
        elif self.color == KEYCODE:
            if len(self.draw_data)  <= 10:
                return Category10[10][rank]
            else:
                return Category20[20][rank % 20]
        return 'black'
    
    def get_label_legend(self, i, info):
        if self.legend == 'In figure':
            return f'{i} - {info[KEYCODE]}'
        return info[KEYCODE]
    
    @param.depends('border', 'x_range_step', 'legend', 'line_width_tree','line_width_chronology', 'circle_radius', 'color', 'axis_marge', watch=True)
    def plot(self, x_range = None, y_range = None):   
        try:
            #print('plot')
            # save x_range and y_range values for next plot (usefull for ligt ploter)
            if x_range is not None:
                self.x_range = x_range
            if y_range is not None:
                self.y_range = y_range
                
            #('ploter')
            self._layout.loading = True
            if self.draw_data is None:
                return
            fig = figure(margin=(5), title=self.figure_title, toolbar_location="left", height=self.height, width=self.width,
                tools="pan,wheel_zoom,box_zoom,reset,hover,save,crosshair", tooltips=[('(date/offset,value)', '(@x, @w)')],
                )
            
            fig.output_backend = "svg"
            radius = self.circle_radius
            
            x = []
            for i, (keycode, info) in enumerate(self.draw_data.items()):
                line_width = self.line_width_tree if info[CATEGORY] == TREE else self.line_width_chronology
                x.append(info['x_min'])
                x.append(info['x_max'])
                fct = fig.line
                if self.param_replot.draw_type == 'Step':
                    fct = fig.step
                info['ids'] = []
                res = fct(x='x', y='y', source=info[HEARTWOOD], line_width=line_width,  color=self.get_color(HEARTWOOD, i), legend_label=self.get_label_legend(i, info))
                #print('plot:', i, res)
                fct(x='x', y='y', source=info[SAPWOOD], line_width=line_width,  color=self.get_color(SAPWOOD, i), legend_label=self.get_label_legend(i, info))
                print('plot cambium', i, info['is_cambium_estimated'])
                if info['is_cambium_estimated'] and self.param_replot.cambium_estimation:
                    fig.segment(x0='x0', y0='y0', x1='x1', y1='y1', source=info[CAMBIUM_ESTIMATED], line_width=line_width, color=self.get_color(SAPWOOD, i), legend_label=self.get_label_legend(i, info))
                    fct(x='x', y='y', source=info[CAMBIUM_BOUNDARIES], line_width=line_width, color=self.get_color(SAPWOOD, i), legend_label=self.get_label_legend(i, info))
                if self.border:
                    fig.quad(top='top', bottom='bottom', left='left', right='right', source=info['border'], line_color='red', alpha=0.05, line_width=1, color='red')
                if MISSING_RING_BEGIN in info:
                    fct(x='x', y='y', source=info[MISSING_RING_BEGIN], line_dash='dotted', line_width=line_width, color=self.get_color(HEARTWOOD, i), legend_label=self.get_label_legend(i, info))
                if MISSING_RING_END in info:
                    c = self.get_color(SAPWOOD, i) if info['is_sapwood'] else self.get_color(HEARTWOOD, i)
                    fct(x='x', y='y', source=info[MISSING_RING_END], line_dash='dotted', line_width=line_width, color=c, legend_label=self.get_label_legend(i, info)) 
                    
                fig.circle(x='x', y='y', source=info[PITH], radius=radius, color=self.get_color(PITH, i), legend_label=self.get_label_legend(i, info))
                fig.circle(x='x', y='y', source=info[CAMBIUM], radius=radius, color=self.get_color(SAPWOOD, i), legend_label=self.get_label_legend(i, info))
                fig.circle(x='x', y='y', source=info[BARK], radius=radius, color=self.get_color(BARK, i), legend_label=self.get_label_legend(i, info))
                
                if self.param_replot.show_dates:
                    fig.add_layout(Label(x=info['x_min'], y=info['y_mean'], text=info[DATE_BEGIN], x_offset=-5 , y_offset=-5, anchor="bottom_right", text_font_size='14px'))
                    fig.add_layout(Label(x=info['x_max'], y=info['y_mean'], text=info[DATE_END], x_offset=5 , y_offset=-5, anchor="bottom_left", text_font_size='14px'))
                
            (x_min, x_max) = (np.min(x), np.max(x)) if self.x_range is None else self.x_range
            fig.x_range = Range1d(start=x_min - self.axis_marge, end=x_max + self.axis_marge)
            if self.y_range is not None:
                fig.y_range = Range1d(self.y_range[0], self.y_range[1])

            fig.legend.visible = False
            self.figure_pane.object = fig
            self.on_x_range_step()
            self.on_legend()
        except Exception as inst:
            logger.error(f'plot : {inst}', exc_info=True)
        finally:
            self._layout.loading = False

