#!/usr/bin/env python3

# -*- coding: utf-8 -*-

from singlecellmultiomics.utils import bdbplot,organoidTools

from copy import deepcopy

from importlib import reload

import numpy as np

class StatePlotter():

    def __init__(self, plot=None):

        self.canvas = bdbplot.BDBPlot() if plot is None else plot

        self.heightPerState = 10

        self.stateMargin = 1

        self.chromosomeMargin = 2

        self.pixelsPerBase = 1/10_000_000

        self.headerHeight=30

        self.bigFontSize = 10

        self.smallFontSize = 6

        self.gray = 220

        self.gainColor = '#B36466'

        self.gainTwoColor = '#FF0087'

        self.normalColor = '#FFFFFF'

        self.lossColor =  '#6562B5'

        self.totalLossColor  = '#004BD8'

        self.missingColor = '#AAAAAA'

    def plotStates(self, df, offset=(0,0), **kwargs):

        self.offset = offset

        stateGroup = self.canvas.getGroup('stateGrid')

        self.stateGroup = stateGroup

        stateGroup.set('transform' ,f"translate({offset[0]},{offset[1]})" )

        self.canvas.svgTree.append(stateGroup)

        self.states = sorted(list(set(df['cluster'])))

        self.chromosomeOrder =  sorted(sorted(list(set(df[ 'chromosome']))), key=organoidTools.chrom_sort_human)

        x = 10

        for i,state in enumerate(self.states):

            g = self.canvas.getGroup(f'state_{state}')

            stateGroup.append(g)

            result = self.plotState(self.canvas, g, df[ df['cluster']==state ], x,

                           y= (self.heightPerState*i + (i*self.stateMargin)), label=state ,

                                   plotChromosomeLabels=(i==0), **kwargs

                                   )

        return self.canvas

    def plotState(self, plot, g, row, x, y, label, plotChromosomeLabels=False,

                  logScale=False, chromosomeSizes=None, logRepeats=False):

        currentX = x

        shownAlleles = False

        for chromosome in self.chromosomeOrder:

            print(chromosome)

            isAllelic = chromosome.endswith('_A') or chromosome.endswith('_B')

            #Obtain how many bases this chromosome has

            chrom = chromosome.split('_')[0]

            if chromosomeSizes is not None and chrom in chromosomeSizes:

                chromosomeSize = chromosomeSizes[chrom]

            else:

                chromosomeSize = row[ row['chromosome']==chromosome ]['endCoordinate'].max()

            chromosomePixelWidth = chromosomeSize * self.pixelsPerBase

            rect =  self.canvas.getRectangle(currentX, y, chromosomePixelWidth, self.heightPerState)

            self.canvas.modifyStyle( rect, {'fill':f'rgb({self.gray},{self.gray},{self.gray})', 'stroke':'none'})

            g.append(rect)

            if plotChromosomeLabels:

                if  not chromosome.endswith('_B'): # allelic

                    text = plot.getText(chromosome.replace('chr','').replace('_A',''), currentX+chromosomePixelWidth*0.5 + (chromosomePixelWidth*0.5 if chromosome.endswith('_A') else 0), y - self.heightPerState)

                    text.set('text-anchor','middle')

                    text.set('dominant-baseline','middle')

                    text.set('font-family','Helvetica')

                    text.set('font-size', str(self.smallFontSize))

                    g.append(text)

                if isAllelic:

                    alleleDescriptor = 'allele ' if not shownAlleles else ''

                    allele = 'A' if chromosome.endswith('_A') else 'B'

                    text = self.canvas.getText( f'{alleleDescriptor}{allele}', currentX+chromosomePixelWidth*0.5 , y - 0.4*self.heightPerState)

                    text.set('text-anchor','middle')

                    text.set('dominant-baseline','middle')

                    text.set('font-family','Helvetica')

                    text.set('font-size', str(self.smallFontSize*0.8))

                    g.append(text)

                    if allele=='B':

                        shownAlleles=True

            if isAllelic:

                offset = self.chromosomeMargin*0.8

                r = self.canvas.getRectangle(currentX-offset, y-self.headerHeight*0.5,

                                      chromosomePixelWidth+offset*2, self.heightPerState+offset*2+self.headerHeight*0.5)

                r.set('z-index', '0')

                self.canvas.modifyStyle( r, {'fill':'#DCFFCE', 'stroke':'none'})

                self.stateGroup.insert(0,r)

            withinX = currentX

            binSizes = {}

            if logScale:

                currCoord = 0

                minBinSize = 1_000

                for binIndex in sorted(list(row[ row['chromosome']==chromosome ]['binIndex'])):

                    d = row[ row['chromosome']==chromosome ]['binIndex']==binIndex

                    dat = row[ row['chromosome']==chromosome ][d].iloc[0,:]

                    cn =dat['copyNumber']

                    space = dat['startCoordinate'] - currCoord

                    if space > minBinSize: # add intermediate bin:

                        binSizes[(binIndex, 'spacer')] = np.log(space)/10 if logRepeats else space/chromosomePixelWidth

                    size = np.log((dat['endCoordinate'] - dat['startCoordinate'] ))

                    binSizes[binIndex] = size

                    currCoord=dat['endCoordinate']

                space = chromosomeSize - currCoord

                if space > minBinSize: # add intermediate bin:

                    binSizes[(binIndex, 'spacerfinal')] = np.log(space)/10 if logRepeats else space/chromosomePixelWidth

            for binIndex in sorted(list(row[ row['chromosome']==chromosome ]['binIndex'])):

                d = row[ row['chromosome']==chromosome ]['binIndex']==binIndex

                dat = row[ row['chromosome']==chromosome ][d].iloc[0,:]

                cn =dat['copyNumber']

                if logScale:

                    if (binIndex, 'spacer') in binSizes:

                        withinX+=( binSizes[(binIndex,'spacer')] / sum( binSizes.values() ) ) * chromosomePixelWidth

                    size = ( binSizes[binIndex] / sum( binSizes.values() ) ) * chromosomePixelWidth

                else:

                    size = (dat['endCoordinate'] - dat['startCoordinate'] )*self.pixelsPerBase

                r = self.canvas.getRectangle(withinX, y, size, self.heightPerState)

                self.canvas.modifyStyle(r, {'stroke-width':'0.2'})

                fillAttr='fill'

                if chromosome.endswith('_A') or chromosome.endswith('_B'): # allelic

                    cn+=1 # (diploid color == white == 2 )

                if cn >= 4:

                    self.canvas.modifyStyle(r, {fillAttr:self.gainTwoColor})

                if cn == 3:

                    self.canvas.modifyStyle(r, {fillAttr:self.gainColor})

                if cn==2:

                    self.canvas.modifyStyle(r, {fillAttr:self.normalColor})

                if cn==1:

                    self.canvas.modifyStyle(r, {fillAttr:self.lossColor})

                if cn==0:

                    self.canvas.modifyStyle(r, {fillAttr:self.totalLossColor})

                if np.isnan(cn):

                    self.canvas.modifyStyle(r, {fillAttr:self.missingColor})

                g.append(r)

                withinX+=size

            currentX += chromosomePixelWidth+self.chromosomeMargin

        ##### plot the label:

        text = self.canvas.getText(label, currentX, y+0.5*self.heightPerState)

        text.set('text-anchor','begin')

        text.set('dominant-baseline','middle')

        text.set('font-family','Helvetica')

        text.set('font-size', str(self.bigFontSize))

        g.append(text)

        self.canvas.setWidth( max(self.canvas.width, currentX+self.chromosomeMargin+20+ self.offset[0] ) )

        self.canvas.setHeight( max(self.canvas.height, y+self.heightPerState*1.5+ self.offset[1] ))

        return {'x':x}