import pyx
from pyx import *

## Set of methods useful in drawing gene organization diagrams

############
## constants
############
X_RANGE = range(1,35) # acceptable range of x-axis parameters
Y_RANGE = range(1,20) # acceptable range of y-axis parameters
DIVR = 5


############
## classes
############
class strand_pt(path.path):
    """stranded divider coming from x1,y2 to x2,y2 with end markers of height endht"""
    def __init__(self,x1,y1,x2,y2,endht):
        path.path.__init__(self, path.moveto_pt(x1,y1+endht+DIVR),
                            path.lineto_pt(x1,y1+DIVR),
                            path.lineto_pt(x1,y1),
                            path.lineto_pt(x2,y2),
                            path.lineto_pt(x2, y2+DIVR),
                            path.lineto_pt(x1,y1+DIVR),
                            path.lineto_pt(x2, y2+DIVR),
                            path.lineto_pt(x2,y2+endht+DIVR),
                            path.lineto_pt(x2,y2-endht),
                            path.lineto_pt(x2,y2),
                            path.lineto_pt(x1,y1),
                            path.lineto_pt(x1,y1-endht))

        
class strand_divider(strand_pt):
    def __init__(self, x1,y1,x2,y2,endht=1):
        """constructor"""
        self.origin = x1,y1
        self.end = x2,y2
        self.mid = abs(self.origin[0] - self.end[0])/2 + self.origin[0]
        strand_pt.__init__(self, unit.topt(x1), unit.topt(y1), unit.topt(x2), unit.topt(y2),unit.topt(endht))
        self.top_top = y1 + endht
        self.top_bottom = y1 + (endht/2.0)
        self.bottom_top = y1 - (endht/2.0)
        self.bottom_bottom = y1 - endht


    def get_placement(self, key):
        """method to retrieve the positional value corrresponding to a strand_divider position
        key -- string key of placement
        returns corresponding y-axis position of key"""
        _placements = { "toptop"        : self.top_top,         # dictionary of valid y-positions in strand_divider
                        "topbottom"     : self.top_bottom,
                        "bottomtop"     : self.bottom_top,
                        "bottombottom"  : self.bottom_bottom }        
        return _placements[key.lower()]


    def strand_desc(self, lbl, dist=2, textattrs=[text.size.tiny], texmsg=[]):
        """prints a string value lbl to the left side of the strand_divider
        lbl -- string value to affix next to this strand_divider
        dist -- distance of the text from the origin
        returns a text object to place"""
        return text.text(self.origin[0] - dist, self.origin[1], lbl, textattrs, texmsg)


    def add_split(self, location, ht=1, position=-1):
        """adds a splitter to the strand divider, to differentiate between sets in fusions and fissions
        location -- "top" or "bottom" to designate where to add the split
        ht -- height of the split
        position -- position of the split, relative to the origin
        returns a line path to draw a split"""
        if position == -1: position = self.mid
        y_ht = self.origin[1]+ht
        if location.lower() == "bottom": y_ht = self.origin[1]-ht
        return pyx.path.line(position+self.origin[0], self.origin[1], position+self.origin[0], y_ht)

    
class arrangement_action:
    """class representing a gene rearrangement"""
    def __init__(self,x,y,action):
        """constructor
        x -- x-axis position of where to locate the action
        y -- y-axis position pertaining to the action location
        type -- string name of the action ('fusion', 'reversal', etc.)
        step -- integer value of the step of the action"""
        self.xpos, self.ypos, self.action = x,y,type,action


############
## functions
############
def approx_len(genelist):
    """approximates the length of a genelist when represented on an EPS file"""
    return .5 + len(genelist) / 10


def genebox(x,y,genelist,label="",textattrs=[text.size.tiny],texmsg=[]):
    """draws a textbox gene box at position (x,y)
    x -- x-axis position of gene
    y -- y-axis position of gene
    label -- label to affix to the gene    
    genelist -- n list of genes in the box (from n = 1...x)
    textattrs -- text attributes
    returns a tuple consisting of the box in the first value, and the text in the second"""
    expr, extr = "", ""
    if str(genelist[0]).find("-") == -1:                # if this is not a complementary gene set (positive values)
        extr = " "
        neg = False
    expr += label + extr + str(genelist[0])
    if len(genelist) > 1:                               # prepare the gene representation        
        expr = label + extr + str(genelist[0]) + ".."   # if the number of genes represented in the box is > 1, then prepare a concatenated list        
        for f in range(len(genelist)/5):            
            expr += "."                                 # for every three genes, draw a '.'        
        expr += str(genelist[len(genelist)-1])          # draw the last gene (to show range)
        expr += extr
    t = text.text(x, y, expr, textattrs, texmsg)
    return t.bbox().enlarged(0.1).path(),t              # return tuple of a formed box and the gene text content


def gene_strand_bind(genelist, str_div, x,ypos,label="",textattrs=[text.size.tiny],texmsg=[]):
    """binds a gene to a strand_divider (also see genebox() for a non-binded gene-set creation)
    gene -- list of genes (in raw unprocessed form of [1,2,3,4,5,...,etc]
    str_div -- strand_divider to bind gene to
    x -- x-axis of where to place the gene on the strand_divider (relative to 0 of str_div)
    y -- y-axis of where to place the gene on the strand_divider
    returns a tuple of formed box and gene text content"""
    return genebox(x+str_div.origin[0],ypos,genelist,label,textattrs,texmsg)
    return genebox(x+str_div.origin[0],ypos,genelist,label,textattrs,texmsg)


def gene_draw_handler(gene, attrs, cnvs):
    """helper function to draw a gene onto a canvas using draw()
    and insert()
    gene -- gene (generated from genebox()) to draw
    attrs -- drawing attributes
    cnvs -- canvas to draw onto"""
    cnvs.draw(gene[0], attrs)   # draw the box with specified attribs
    cnvs.insert(gene[1])        # draw the gene numbers


def path_draw_handler(path_tuple, attrs, cnvs):
    """helper function to draw a set of path canvas.draw() based tuples
    path_tuples -- tuple of paths to draw
    attrs -- canvas.draw() attributes
    cnvs -- canvas to draw onto"""
    try:
        len(path_tuple)
    except AttributeError:              # test if path_tuple is not a tuple
        cnvs.draw(path_tuple, attrs)
        return
    for p in path_tuple:                # if path_tuple == tuple, draw each path
        cnvs.draw(p, attrs)


def simple_reversal(gensrc, gentgt, boxdist=.2):
    """prepares a path curve with only one line/arrow going from gensrc to gentgt
    gensrc -- source set of genes
    gentgt -- target genes(e.g. arrow points at this gene)
    boxdist -- distance of the line from the gene markers
    returns a tuple consisting of the arcs to stroke onto the canvas"""    
    return pyx.connector.line(gensrc, gentgt, boxdists=boxdist)


def reversal(genset1, genset2, boxdist=.2, angle=45):
    """prepares a tuple of arcs that represent a reversal action on the set of genes passed, where a reversal is
    a conserved list-wise inversion of a set of genes
    genset1 -- first set of genes, source
    getnset2 -- second set of genes, target (e.g. arrow points at this gene)
    boxdist -- distance of the line from the gene markers
    angle -- angle of exit from the source (one arc will be at angle, the other at -angle)
    returns a tuple consisting of the arcs to stroke onto the canvas"""
    return  pyx.connector.arc(genset1, genset2, boxdists=boxdist, relangle=angle), pyx.connector.arc(genset1, genset2, boxdists=boxdist, relangle=-angle)


def fission(gensrc, gentgt1, gentgt2, boxdist=.2):
    """prepares a tuple of line paths that represent a fission action on a set of genes passed, where a fission is
    the division of 1 gene set into 2 (and only 2) separate gene sets
    gensrc -- source gene set
    gentgt1 -- first set of genes resultant from fission
    gentgt2 -- second set of genes resultant from fission
    boxdist -- distance of the path from the gene sets
    returns a tuple consisting of two pathes, each one origininating from gensrc and ending at either gentgt1 or gentgt2"""
    return  pyx.connector.line(gensrc, gentgt1, boxdists=boxdist), pyx.connector.line(gensrc, gentgt2, boxdists=boxdist)


def fusion(gentgt, gensrc1, gensrc2, boxdist=.2):
    """prepares a tuple of line paths that represent a fusion action on a set of genes passed, where a fusion is the
    combination of 2 (and only 2) sets of genes into a single set of genes
    gentgt -- resultant set of genes from fusion
    gensrc1 -- first gene set source for fusion
    gensrc2 -- second gene set source fro fusion
    boxdist -- distance of the path from the gene sets
    returns a tuple consisting of two paths, either originating from a gensrc, and ending at gentgt"""
    return  pyx.connector.line(gensrc1, gentgt, boxdists=boxdist), pyx.connector.line(gensrc2, gentgt, boxdists=boxdist)


def translocation(genset1, genset2, boxdist=.2):
    """prepares a tuple of line paths that represent a translocation action on a set of genes passed, where a
    translocation is the conserved swapping of positions between two different sets of genes
    genset1 -- first gene set to swap
    genset2 -- second gene set to swap
    boxdist -- distance of the path from the gene sets
    returns a tuple consisting of two paths, one coming going from genset1 to genset2, and the other vice versa"""
    return  pyx.connector.line(genset1, genset2, boxdists=boxdist), pyx.connector.line(genset2, genset1, boxdists=boxdist)


############
## testing
############
if __name__ == '__main__':

    c = canvas.canvas()

    # example action
    sd = strand_divider(X_RANGE[0],10,X_RANGE[len(X_RANGE)-1],10)    
    c.stroke(sd,[style.linewidth.Thin])
    c.stroke(sd.add_split("top"),[style.linestyle.dashed, style.linewidth.Thin])
    c.stroke(sd.add_split("bottom", position=5),[style.linestyle.dashed, style.linewidth.Thin])
    c.insert(sd.strand_desc("999-Translocation"))

    # example set for reversal
    #gb1 = genebox(2, 20,[1,2,3,4,5])
    #gb2 = genebox(2, 1,[-5,-4,-3,-2,-1])    
    gb1 = gene_strand_bind([-234, -235, -236, -237], sd, 6, sd.get_placement("toptop"))
    gb2 = gene_strand_bind([-5,-4,-3,-2,-1], sd, 6, sd.get_placement("bottombottom"))
    
    # example set for fusion    
    #gba = genebox(7.5,7,[1,2,3,4,5])
    #gbb = genebox(9,7,[6,7,8,9,10])
    #gbc = genebox(8,6,[1,2,3,4,5,6,7,8,9,10])
    gba = gene_strand_bind([1,2,3,4,5], sd, 8, sd.get_placement("toptop"))
    gbb = gene_strand_bind([6,7,8,9,10], sd, 9, sd.get_placement("topbottom"))
    gbc = gene_strand_bind([1,2,3,4,5,6,7,8,9,10], sd, 8, sd.get_placement("bottomtop"))

    # example set for fission
    #gbx = genebox(7.5, 5, [1,2,3,4,5,6,7,8,9,10])
    #gby = genebox(8, 4, [1,2,3,4,5])
    #gbz = genebox(9.5, 4, [6,7,8,9,10])
    gbx = gene_strand_bind([1,2,3,4,5,6,7,8,9,10], sd, 4, sd.get_placement("topbottom")) 
    gby = gene_strand_bind([1,2,3,4,5], sd, 3, sd.get_placement("bottombottom")) 
    gbz = gene_strand_bind([6,7,8,9,10], sd, 4, sd.get_placement("bottomtop")) 

    # example for translocation
    #gb9 = genebox(25, 7.2, [1,2,3,4,5])
    #gb8 = genebox(25, 6, [6,7,8,9,10])
    gb9 =  gene_strand_bind([1,2,3,4,5], sd, 10, sd.get_placement("topbottom"))
    gb8 =  gene_strand_bind([6,7,8,9,10], sd, 10, sd.get_placement("bottomtop")) 

    #reversal code    
    #rev = reversal(gb1[1],gb2[1],angle=45)
    rev = simple_reversal(gb1[1],gb2[1])
    gene_draw_handler(gb1, [deco.stroked(), deco.filled([color.grey(0.85)])], c)
    gene_draw_handler(gb2, [deco.stroked(), deco.filled([color.rgb.blue])], c)
    path_draw_handler(rev, [color.rgb.blue, deco.earrow.normal], c) # draw paths

    # fusion code
    fus = fusion(gbc[1], gba[1], gbb[1])
    gene_draw_handler(gba, [deco.stroked(), deco.filled([color.grey(0.85)])], c)
    gene_draw_handler(gbb, [deco.stroked(), deco.filled([color.grey(0.85)])], c)
    gene_draw_handler(gbc, [deco.stroked(), deco.filled([color.rgb.red])], c)
    path_draw_handler(fus, [color.rgb.red, deco.earrow.normal], c) # draw paths

    # fission code
    fis = fission(gbx[1], gby[1], gbz[1])
    gene_draw_handler(gbx, [deco.stroked(), deco.filled([color.grey(0.85)])], c)
    gene_draw_handler(gby, [deco.stroked(), deco.filled([color.rgb.green])], c)
    gene_draw_handler(gbz, [deco.stroked(), deco.filled([color.rgb.green])], c)
    path_draw_handler(fis, [color.rgb.green, deco.earrow.normal], c) # draw paths

    # translocation code
    trans = translocation(gb9[1], gb8[1])
    gene_draw_handler(gb9, [deco.stroked(), deco.filled([color.grey(0.5)])], c)
    gene_draw_handler(gb8, [deco.stroked(), deco.filled([color.grey(0.5)])], c)
    path_draw_handler(trans, [color.grey(0.5), deco.earrow.normal], c) # draw paths

    c.writeEPSfile("test", paperformat='LETTER', fittosize=True,rotated=90)