#!/usr/bin/python3

import numpy
import matplotlib.pyplot

matplotlib.pyplot.ion()

global scale
scale = 1.0

class Level:
    def __init__(self, indx, energy, start, stop, energylabel="", spinlabel="", color='k'):
        self.indx = indx
        self.energy = energy
        self.start = start
        self.stop = stop
        self.energylabel = energylabel
        self.spinlabel = spinlabel
        self.color = color
        self.diag = 0.0
        self.extension = False
        
    def add_extension(self, start, stop):
        self.extstart = start
        self.extstop = stop
        self.extension = True
        
    def draw(self, ax, col='k', ls='-'):
        if (self.extension):
            ax.plot([self.extstart, self.extstop], [self.energy, self.energy], color='grey', linestyle=':', zorder=1)

        ax.plot([self.start, self.stop], [self.energy, self.energy], color=self.color, linestyle=ls, zorder=2)
        ax.text(self.start, self.energy, self.energylabel, ha='left', va='bottom')
        ax.text(self.stop, self.energy, self.spinlabel, ha='right', va='bottom')
        if (self.indx >= 0 and self.diag != 0.0):
            ax.text(self.stop, self.energy-10, '%3.1f'%(self.diag), ha='right', va='top')
        
class Transition:
    def __init__(self, indi, indf, start, stop, be2, headlength=3.0, color='k'):
        global scale
        self.indi = indi
        self.indf = indf
        self.start = numpy.array(start)
        self.stop = numpy.array(stop)
        self.be2 = be2
        self.alpha = 1.0
        #self.width = max(be2*scale, 100*scale)
        self.headlength = headlength
        self.color = color
        
    def setBE2(self, be2):
        global scale
        self.be2 = be2
        self.width = be2*scale
        
    def getwidth(self):
        global scale
        return min(100*scale, max(self.be2*scale, 5.0*scale))
    
    def getlength(self):
        length = numpy.sqrt((self.stop-self.start)[0]**2 + (self.stop-self.start)[1]**2)
        return length
    
    def getheadwidth(self):
        return max(min(self.getwidth()*1.5, self.getwidth() + 50), self.getwidth() + 5)
    
    def getheadlength(self):
        return min(self.headlength*self.getheadwidth(), self.getlength()/2.0) 
        
    def draw(self, ax):
        self.arrow(ax)
        mid = (self.stop + self.start)/2.0
        props = dict(boxstyle='square', facecolor='white', alpha=0.7, lw=0, zorder=11)
        ax.text(mid[0], mid[1],  '{:g}'.format(float('{:.{p}g}'.format(self.be2, p=2))), bbox=props, ha='center', va='center', zorder=12)
        
    def arrow(self, ax):
        x = []
        y = []
        end = self.stop
        start = self.start
        direction = end-start
        length = numpy.sqrt((end-start)[0]**2 + (end-start)[1]**2)
        direction /= length
        disp_dir = ax.transData.transform(end) - ax.transData.transform(start)
        disp_length = numpy.sqrt(disp_dir[0]**2 + disp_dir[1]**2)
        disp_dir /= disp_length
        perp_dir = numpy.array((-disp_dir[1], disp_dir[0]))
        disp_start = ax.transData.transform(start)

        perp_ang = numpy.arctan2(perp_dir[1], perp_dir[0])
        topwidth = (self.getwidth()/2.0)/numpy.cos(perp_ang)

        start1 = ax.transData.inverted().transform(disp_start - numpy.array([topwidth, 0]))
        start2 = ax.transData.inverted().transform(disp_start + numpy.array([topwidth,0]))
        x.append(start1[0])
        y.append(start1[1])
        x.append(start2[0])
        y.append(start2[1])

        headpos = start + direction*(length-self.getheadlength())
        disp_headpos = numpy.array(ax.transData.transform(headpos))

        pos1 = ax.transData.inverted().transform(disp_headpos + self.getwidth()/2.0*perp_dir)
        pos2 = ax.transData.inverted().transform(disp_headpos + self.getheadwidth()/2.0*perp_dir)
        pos3 = ax.transData.inverted().transform(disp_headpos - self.getheadwidth()/2.0*perp_dir)
        pos4 = ax.transData.inverted().transform(disp_headpos - self.getwidth()/2.0*perp_dir)
        x.append(pos1[0])
        y.append(pos1[1])
        x.append(pos2[0])
        y.append(pos2[1])
        x.append(end[0])
        y.append(end[1])
        x.append(pos3[0])
        y.append(pos3[1])
        x.append(pos4[0])
        y.append(pos4[1])
        
        ax.fill(x,y,color=self.color,alpha=self.alpha,zorder=10)
        #ax.plot(x,y,color)

class Scheme:
    def __init__(self, filename, colorfilename):
        file = open(filename)
        self.levels = {}
        self.transitions = {}
        self.colors = {}
        colorfile = open(colorfilename)
                
        for l in colorfile.readlines():
            l = l.split()
            self.colors[int(l[0])] = (float(l[1]), float(l[2]), float(l[3]))
        
        for l in file.readlines():
            l = l.split()
            if l[0] == "#":
                continue
            if l[0] == "L":
                self.levels[int(l[1])] = Level(int(l[1]), float(l[2]), float(l[3]), float(l[4]), l[5], "$"+l[6]+"$",
                                              color='k')                
            elif l[0] == "T":
                start = self.levels[int(l[1])].start + (self.levels[int(l[1])].stop - self.levels[int(l[1])].start)*float(l[4])
                stop = self.levels[int(l[2])].start + (self.levels[int(l[2])].stop - self.levels[int(l[2])].start)*float(l[5])

                self.transitions[(int(l[1]), int(l[2]))] = Transition(int(l[1]), int(l[2]), 
                                                                      [start, self.levels[int(l[1])].energy],
                                                                      [stop, self.levels[int(l[2])].energy], 
                                                                      float(l[3]),
                                                                     color=self.colors[int(l[6])])
            elif l[0] == "E":
                self.levels[int(l[1])].add_extension(float(l[2]), float(l[3]))
    
            
    
    def draw(self, ax):
        for ind, l in self.levels.items():
            l.draw(ax)
    
        for ind, t in self.transitions.items():
            t.draw(ax)


fig,ax = matplotlib.pyplot.subplots();
scheme = Scheme("prm_levels.scheme", "./ROOT_colorwheel.dat")

scheme.draw(ax)

ax.set_xlabel("Dominant SP configuration")
ax.set_ylabel("Energy (keV)")

matplotlib.pyplot.show(block=True);
