from biogeneutil import *
from grimm import *
from pyx import *
from genepyx import *
from pyx.connector import line
import random
import genepyx
import string

DECR_Y = 6
CANVAS_MAX = 20
BIND_ADJ = 3.0

## Classes and functions that allow for the graphing of genome rearrangement scenarios
############
## Functions
############
def approx_gene_len(genelist):
    """approximates the length of a genelist when represented on an EPS file"""
    return 1.7 + len(genelist)/35
    

############
## Classes
############
class ColorManager:
    """class that handles the doling out of colors to various denotions"""    
    clist = [ color.cmyk.Dandelion, color.cmyk.PineGreen, color.cmyk.Salmon, \
               color.cmyk.Periwinkle,  color.cmyk.GreenYellow,   color.cmyk.CadetBlue,\
               color.cmyk.CornflowerBlue,  color.cmyk.Cerulean,  color.cmyk.Aquamarine,\
               color.cmyk.Green, color.cmyk.Thistle, color.cmyk.SpringGreen,\
               color.cmyk.Tan,  color.cmyk.Gray ]


    def __init__(self, repeat=True, colorlist=clist):
        """constructor - initializes the color dictionary
        colorref -- default color base to use"""
        self.color_list = colorlist
        random.shuffle(self.color_list)
        self.use_repeat = repeat
        self.color_cursor = 0
        self.color_ref = {}


    def colorize(self, obj):
        """colorizes based on the object passed
        obj -- object to colorize        
        returns a color based on the object"""
        try:
            return self.color_ref[obj]      # return the proper color if this obj has been added before
        except KeyError:                    # otherwise, give it a new color, save, and return that color
            if not self.use_repeat: assert self.color_cursor < len(self.color_list), 'Out of colors'
            elif self.color_cursor >= len(self.color_list): self.color_cursor = 0
            self.color_ref[obj] = self.color_list[self.color_cursor]            
            self.color_cursor += 1
            return self.color_ref[obj]


    def getkeys(self):
        """returns the key (legend) in a list tuple form, with the first element being the value,
        and the second element being the color"""
        return [(self.color_ref[y], y) for y in self.color_ref.keys()]        


class DrawHandler:
    """class to handle drawing the rearrangement network"""
    def __init__(self, drawablelist, initfilename, tcm=ColorManager()):
        """DrawHandler constructor
        drawablelist -- list of RearrangementWrappers (with GRIMMActions) to draw
        initfilename -- initial filename to write the EPS to"""
        self.workset, self.filename = drawablelist, initfilename
        self.sdlist = []    # list of strand_dividers
        self.cnvslst = []   # list to hold the canvas objects for drawing
        self.cm = tcm # color manager to manage the colors used
        cnvscnt, ycnt = 0, 25
        curcanv = canvas.canvas()
#        print len(drawablelist)
        for r in drawablelist:
            ttlsize = 0
            for t in r.get_grimmaction().get_before(): ttlsize += len(consec_list(t))    # get size estimates to gauge the length of the strand
            for t in r.get_grimmaction().get_after(): ttlsize += len(consec_list(t))            
            sd = strand_divider(0,ycnt,ttlsize,ycnt)
            ycnt -= DECR_Y      # decrement the y-position (move down to the next strand) and            
            cnvscnt += 1        # increment the count of strands on the canvas (at the CANVAS_MAX,
            self._draw_strand_to_canvas(sd, curcanv, r)
            self._process_states(r, sd, curcanv)        # process the actions onto the strand_divider on canvas
            if not cnvscnt < CANVAS_MAX:                # go to a new canvas)
                self._draw_legend(curcanv, ycnt)
                self.cnvslst.append(curcanv)
                curcanv = canvas.canvas()
                ycnt, cnvscnt = 35, 0
        self._draw_legend(curcanv, ycnt)
        if self.cnvslst == [] or not curcanv is self.cnvslst[-1]: self.cnvslst.append(curcanv) # add the last set (remainder)        
        self._draw_canvases(initfilename)


    def _draw_legend(self, cnv, ypos):
        """draws the legend based on the ColorManager(cm)
        cnv -- canvas to draw onto
        ypos -- y-axis position on the canvas to draw to"""
        keyd = self.cm.getkeys()
        for tup in keyd:
            cnv.stroke(path.rect(3, ypos, 1, .5), [deco.filled([tup[0]])])
            val =  str(tup[1])
            val = string.replace(val, "_", "")  # TeXRunner has problems with underscore
            val = val[:-2]            
            cnv.text(4.2, ypos, val)
            ypos -= 1


    def _draw_strand_to_canvas(self, strand, cnv, action):
        """draws the strand onto canvas
        strand -- strand to draw
        cnv -- canvas to draw
        action -- associated action (in wrapper)"""
        step = action.get_grimmaction().get_action().lower()
        desc = action.get_grimmaction().get_step() + "-" + step
        cnv.stroke(strand,[style.linewidth.thick])       # draw the strand_divider
        cnv.insert(strand.strand_desc(desc))            # draw strand


    def _process_states(self, action, strand, canv):
        """processes the before and after states of action for the respective strand
        onto the canvas canv
        action -- RearrangementWrapper to process get_before() and get_after() of
        strand -- strand_divider to add actions to
        canv -- canvas.canvas() to use"""
        if action.get_grimmaction().get_action() == 'reversal': self._process_reversal(action, strand, canv)
        elif action.get_grimmaction().get_action() == 'translocation': self._process_translocation(action, strand, canv)
        elif action.get_grimmaction().get_action() == 'fusion': self._process_fusion(action, strand, canv)
        elif action.get_grimmaction().get_action() == 'fission': self._process_fission(action, strand, canv)


    def _process_handler(self, action_altered, action_state, strand, cnv, position, start=.5, dictionary=None):
        """method that handles calling the binding function for all action processes
        action_state -- action state (before, after) list with data to bind
        action_altered -- action altered array of action (bad design, I know)
        strand -- strand_divider to use
        cnv -- canvas to bind to
        position -- 'top' or 'bottom'
        returns a tuple consisting of the binded genes, the original gene list to bind, and
        the last binding position on the strand"""
        a_marked = consec_list(action_altered)                                      # consecutize the altered genes
        c_marked = adv_discrim_list(action_state, a_marked)                         # divide the gene set by the altered genes and conesecutives
        a_binded = self._init_bind_all(c_marked, strand, cnv, position, a_marked, start, dictionary)   # bind to canvas
        return a_binded
    
            
    def _process_reversal(self, action, strand, canv):
        """specifically processes reversal actions by binding relevant genes to the given strand
        action -- RearrangementWrapper to process get_before() and get_after() of
        strand -- strand_divider to add actions to
        canv -- canvas.canvas() to use"""
        self._process_handler(action.get_grimmaction().get_before_altered()[0], action.get_grimmaction().get_before()[0], strand, canv, "top",dictionary=action.get_dictionary())
        self._process_handler(action.get_grimmaction().get_after_altered()[0], action.get_grimmaction().get_after()[0], strand, canv, "bottom",dictionary=action.get_dictionary())


    def _process_translocation(self, action, strand, canv):
        """specifically processes transloc actions by binding relevant genes to the given strand
        action -- RearrangementWrapper to process get_before() and get_after() of
        strand -- strand_divider to add actions to
        canv -- canvas.canvas() to use"""
        binded = self._process_handler(action.get_grimmaction().get_before_altered()[0], action.get_grimmaction().get_before()[0], strand, canv, "top",dictionary=action.get_dictionary())
        canv.stroke(strand.add_split("top",position=binded[2]+BIND_ADJ/2),[style.linestyle.dashed, style.linewidth.thick])
        self._process_handler(action.get_grimmaction().get_before_altered()[1], action.get_grimmaction().get_before()[1], strand, canv, "top",start=binded[2]+BIND_ADJ,dictionary=action.get_dictionary())

        binded = self._process_handler(action.get_grimmaction().get_after_altered()[0], action.get_grimmaction().get_after()[0], strand, canv, "bottom",dictionary=action.get_dictionary())
        canv.stroke(strand.add_split("bottom",position=binded[2]+BIND_ADJ/2),[style.linestyle.dashed, style.linewidth.thick])
        self._process_handler(action.get_grimmaction().get_after_altered()[1], action.get_grimmaction().get_after()[1], strand, canv, "bottom",start=binded[2]+BIND_ADJ,dictionary=action.get_dictionary())


    def _process_fission(self, action, strand, canv):
        """specifically processes transloc actions by binding relevant genes to the given strand
        action -- RearrangementWrapper to process get_before() and get_after() of
        strand -- strand_divider to add actions to
        canv -- canvas.canvas() to use"""
        # process one side of the before scenario
        try:
            self._process_handler(action.get_grimmaction().get_before_altered()[0], action.get_grimmaction().get_before()[0], strand, canv, "top",dictionary=action.get_dictionary())
        except ValueError: pass
        # process second side of the before scenario
        try:
            self._process_handler(action.get_grimmaction().get_before_altered()[1], action.get_grimmaction().get_before()[1], strand, canv, "top",dictionary=action.get_dictionary())
        except ValueError: pass
        
        binded = self._process_handler(action.get_grimmaction().get_after_altered()[0], action.get_grimmaction().get_after()[0], strand, canv, "bottom",dictionary=action.get_dictionary())
        canv.stroke(strand.add_split("bottom",position=binded[2]+BIND_ADJ/2),[style.linestyle.dashed, style.linewidth.thick])
        self._process_handler(action.get_grimmaction().get_after_altered()[1], action.get_grimmaction().get_after()[1], strand, canv, "bottom", start=binded[2]+BIND_ADJ,dictionary=action.get_dictionary())


    def _process_fusion(self, action, strand, canv):
        """specifically processes transloc actions by binding relevant genes to the given strand
        action -- RearrangementWrapper to process get_before() and get_after() of
        strand -- strand_divider to add actions to
        canv -- canvas.canvas() to use"""
        # process one side of the before scenario
        binded = self._process_handler(action.get_grimmaction().get_before_altered()[0], action.get_grimmaction().get_before()[0], strand, canv, "top",dictionary=action.get_dictionary())
        canv.stroke(strand.add_split("top",position=binded[2]+BIND_ADJ/2),[style.linestyle.dashed, style.linewidth.thick])
        self._process_handler(action.get_grimmaction().get_before_altered()[1], action.get_grimmaction().get_before()[1], strand, canv, "top", start=binded[2]+BIND_ADJ,dictionary=action.get_dictionary())
        
        try:
            self._process_handler(action.get_grimmaction().get_after_altered()[0], action.get_grimmaction().get_after()[0], strand, canv, "bottom",dictionary=action.get_dictionary())
        except ValueError: pass
        # process second side
        try:
            self._process_handler(action.get_grimmaction().get_after_altered()[1], action.get_grimmaction().get_after()[1], strand, canv, "bottom",dictionary=action.get_dictionary())
        except ValueError: pass


    def _draw_canvases(self, fname):
        """draws all canvases onto multiple EPS files with the initial name fname"""
        cnt = 0
        dname = fname
        for x in self.cnvslst:
            try:
                x.writeEPSfile(dname + str(cnt), paperformat='LETTER', fittosize=True)
            except AttributeError:  # BUG: find out why I need this later
                pass
            cnt += 1


    def _init_bind_all(self, bindablelist, sd, cnvs, whichstrand, marked=None, start_pos=.5, dictionary=None):
        """binds all genes in bindablelist to strand_divider sd; assumes the binding begins at the
        origin of sd; the bindable list is a list taken from a before_action() or after_action() of a
        GRIMMAction, and must be turned into a consec_list() (see biogeneuntil).        
        bindablelist -- list of a list of genes in integer form (signed) to bind to sd
        sd -- strand_divider object to bind to
        cnvs -- canvas to draw to
        whichstrand -- indicates which strand to bind to - "top" for top, "bottom" for bottom
        marked -- list of gene sets to mark (make visible)
        start_pos -- starting positing of writing on the strand_divider
        dictionary -- RearrangementWrapper dictionary of genes (to write out labels)
        returns a tuple of the list of binded genes to sd via genepyx.gene_strand_bind(), the original
        gene list (order between the binded genes and the gene list is perserved), and the last
        position of the binding cursor"""
        bindpos = start_pos
        binded = []
        whichstrand = whichstrand.lower()
        prevgene = None
        for x in bindablelist:                                          # iterate through each geneset in the list            
            curcolor = [self.cm.colorize(dictionary[abs(x[0])][0])]       # retrieve the appropriate color based on the source genome
            fill = deco.filled(curcolor)
            attrs = [deco.stroked(), fill, style.linewidth.Thin, color.cmyk.White]
            if not marked == None:                                      # if marking, then make the altered genes yellow
                #if x in marked: fill = deco.filled([color.cmyk.Yellow])
                if x in marked: attrs = [deco.stroked(), fill, color.cmyk.Black, style.linewidth.THick]
            thestrand = whichstrand                                     # add the proper designation to the strand determinator
            if x[0] > 0: thestrand += "top"
            else: thestrand += "bottom"                
            gene = gene_strand_bind(x, sd, bindpos, sd.get_placement(thestrand))    # make the binding
            
            mrkcol = color.cmyk.Black
            if not marked == None:
                if x in marked: mrkcol = color.cmyk.Red                

            if not prevgene == None: cnvs.stroke(line(prevgene, gene[1], boxdists=0.3), [mrkcol, style.linewidth.Thick])
            prevgene = gene[1]

            binded.append(gene)                                                     # save binding to a list (for use later, maybe? I dunno)
            gene_draw_handler(gene, attrs, cnvs)                   # draw the gene
            if not dictionary == None:              # if using dictionary, write out the chrcontigs to the genes
                listdict = {}
                for t in x:                         # must associate every gene in the geneset x with a chrcontig
                    listdict[dictionary[abs(t)]] = 0
                if whichstrand == "bottom": decr = -.25
                else: decr = .25
                cury = sd.get_placement(thestrand) + decr
                for key in listdict:
                    val =  str(key[0])
                    val = string.replace(val, "_", "")  # TeXRunner has problems with underscore
                    val = val[:-2]
                    cnvs.text(bindpos, cury, val, textattrs=[text.size.tiny, mrkcol])
                    cury += decr
                    val =  str(key[1])
                    val = string.replace(val, "_", "")  # TeXRunner has problems with underscore
                    val = val[:-2]
                    cnvs.text(bindpos, cury, val, textattrs=[text.size.tiny, mrkcol])
                    cury += decr
            bindpos += approx_gene_len(x)
        return binded, bindablelist, bindpos


###########
## testing
###########
if __name__ == '__main__':
    genes = open_grimm('c:\\Python23\\pylib\\tb.txt', '01_01') # save the genes from Tb02_01
    f = locate('c:\\Python23\\pylib\\lmjf.txt', genes) # find the corresponding chr/contigs that contain genes from Tb02_01
    filehandle = open('c:\\Python23\\pylib\\index.html', 'r') # open the GRIMM results
    parser = GRIMMScraper(475)
    while True: # read in the HTML data
        data = filehandle.read(10000)
        if not data: break
        parser.feed(data)
    actionlist = parser.get_grimm_actions() # compile the GRIMMAction list
    processedfile = parser.close()
    filehandle.close()
    relevants = get_relevant_actions(genes,actionlist) # get all the GRIMMActions that are relevant for Tb_02_01
    lmjf = load_chrcontig('c:\\Python23\\pylib\\lmjf.txt', [])
    tb = load_chrcontig('c:\\Python23\\pylib\\tb.txt', [])
    # build the RearrangementNetwork
    network = RearrangementNetwork(actionlist, tb, lmjf)
    
#    h = network.trace_intersect_chromosome('Lmjf_25_01','Tb_11_03')
    thecm = ColorManager()

    h = network.trace_chromosome('')
    d = DrawHandler(h, 'Total_', thecm)