# -*- encoding: utf-8 -*-
'''
@File        :part.py
@Time        :2021/07/11 12:44:00
@Author        :charles kiko
@Version        :1.0
@Contact        :charles_kiko@163.com
@Desc        :家族分布染色体
'''

import csv
import sys
import re
from math import *
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import *
from matplotlib.patches import Circle, Ellipse, Arc
from pylab import *
from famCircle.bez import *


class part():
    def __init__(self, options):
        self.makers = 200
        for k, v in options:
            setattr(self, str(k), v)
            print(k, ' = ', v)


    def plot_bez_inner(self, ex1x, ex1y, ex2x, ex2y):
        x = [ex1x, ex1x+((ex2x - ex1x)/3+0.1), ex2x-((ex2x - ex1x)/3+0.1), ex2x]
        y = [ex1y, ex1y+((ex2y - ex1y)/3), ex2y-((ex2y - ex1y)/3), ex2y]
        step = .01
        t = arange(0, 1+step, step)
        xt = self.Bezier(x, t)# 贝塞尔曲线
        yt = self.Bezier(y, t)
        plot(xt, yt, '-', color='r', lw=.3, alpha=0.3)#alpha 透明度

    def calculate_coef(self,p0, p1, p2, p3):
        c = 3*(p1 - p0)
        b = 3*(p2 - p1) -c
        a = p3 - p0 - c - b
        return c, b, a
    def Bezier(self,plist, t):
        p0, p1, p2, p3 = plist
        # calculates the coefficient values
        c, b, a = self.calculate_coef(p0, p1, p2, p3)
        tsquared = t**2
        tcubic = tsquared*t
        return a*tcubic + b*tsquared + c*t + p0

    def make_plot(self,lens,gff,family,blast,ks,family_dic):
        chro_lingth = lens[self.chro_name]['end']
        chro_order = lens[self.chro_name]['order']
        gene_average = chro_lingth/chro_order
        maker = 0.2
        fig1 = plt.figure(num=1, figsize=(10, 10))  # 确保正方形在屏幕上显示一致，固定figure的长宽相等
        axes1 = fig1.add_subplot(1, 1, 1)
        plt.xlim((0, 1))
        plt.ylim((0, 1))
        self.map1(axes1,maker,0.1,0.3,0.8,'k')# 长染色体0.8
        plt.text(0.45, 0.28, self.chro_name)
        plt.axis('off')
        self.pair_index(axes1,lens,gff,family,blast,ks,gene_average,chro_lingth,chro_order,family_dic)
        plt.savefig(self.savefile,dpi=1000)

    def map1(self,axes1,maker,x,y,x1,color):
        # makers = 15 #标尺
        # x = 0.1 #起始坐标x
        # y = 0.1 #起始坐标y
        # x1 = 0.8#染色体长度
        # color='k'
        lw = 1  #比例线宽
        y1 = 0.001
        w = 0.01
        alpha = 1
        # print(y+w+y1, x, x+x1,y, x, x+x1, lw)
        plt.axhline(y=y, xmin=x, xmax=x+x1, lw=lw, c=color, alpha=alpha)
        plt.axhline(y=y+w+y1, xmin=x, xmax=x+x1, lw=lw, c=color, alpha=alpha)
        for i in range(int(0.8/maker)+1):
            mx = maker*i
            plt.axvline(x=0.1+mx, ymin=y, ymax=y+(w/2.5), lw=lw, c=color, alpha=alpha)
        base1 = Arc(xy=(x, y+((w+y1)/2)),    # 椭圆中心，（圆弧是椭圆的一部分而已）
                width=w+y1,    # 长半轴
                height=w+y1,    # 短半轴
                angle=90,    # 椭圆旋转角度（逆时针） 
                theta1=0,    # 圆弧的起点处角度
                theta2=180,    # 圆度的终点处角度
                color=color,
                alpha=alpha,
                linewidth=lw
                )
        base2 = Arc(xy=(x1+x, y+((w+y1)/2)),    # 椭圆中心，（圆弧是椭圆的一部分而已）
                width=w+y1,    # 长半轴
                height=w+y1,    # 短半轴
                angle=-90,    # 椭圆旋转角度（逆时针） 
                theta1=0,    # 圆弧的起点处角度
                theta2=180,    # 圆度的终点处角度
                color=color,
                alpha=alpha,
                linewidth=lw   #线宽像素
                )
        axes1.add_patch(base1)
        axes1.add_patch(base2)

    def map2(self,axes1,x,y,x1,color):
        # makers = 15 #标尺
        # x = 0.1 #起始坐标x
        # y = 0.1 #起始坐标y
        # x1 = 0.8#染色体长度
        # color='k'
        lw = 2  #比例线宽
        alpha = 1
        plt.axhline(y=y, xmin=x, xmax=x+x1, lw=lw, c=color, alpha=alpha)

    def transform_pt2(self, rad, r):
        return r*cos(rad), r*sin(rad)

    def plot_bez_Ks2(self, ex1x, ex1y, ex2x, ex2y, col, ratio):
        sita = pi / 2
        if ex1x != ex2x:
            sita = atan((ex2y-ex1y)/(ex2x-ex1x))
        d = sqrt((ex2x-ex1x)**2+(ex2y-ex1y)**2)
        L = d * ratio
        P1x = ex1x + L*sin(sita)
        P1y = ex1y - L*cos(sita)
        P2x = ex2x + L*sin(sita)
        P2y = ex2y - L*cos(sita)
        step = .01
        t = arange(0, 1+step, step)
        x=[ex1x, P1x, P2x, ex2x]
        y=[ex1y, P1y, P2y, ex2y]
        # print('x,y,t',x,y,t)
        xt = Bezier(x,t)
        yt = Bezier(y,t)
        # xx3 = xt[int(len(xt)/2)]
        # yy3 = yt[int(len(yt)/2)]-0.001
        # s = '<-'+str(gene1)+'-'+ str(gene2)+'->'
        # text(xx3,yy3, s = str(s),color=col, fontsize = 1.5, rotation=30, alpha=0.8)# 添加基因标注
        plot(xt, yt, '-', color = col, lw = 0.3)#0.1

    def cluster0(self,cluster,id1,id2):
        for i in cluster:
            if id1 in i and id2 in i:
                return True
            else:
                continue
        return False

    def clusterx(self,j,lt):
        for i in range(1,int(self.cluster)):
            if j+i in lt:
                return True
        return False

    def pair_index(self,axes1,lens,gff,family,blast,ks,gene_average,chro_lingth,chro_order,family_dic):
        fam_dic = {}
        for i in family:
            if gff[i]['chr'] != self.chro_name:
                continue
            else:
                fam_dic[gff[i]['order']] = i

        # cluster = []# 基因簇
        lt = list(fam_dic.keys())
        lt.sort()# 升序
        # print(lt)
        cluster = []
        cluster0 = []
        for j in lt:
            if self.clusterx(j,lt):
                cluster0.append(fam_dic[j])
            else:
                cluster.append(cluster0)
                cluster0 = []
                cluster0.append(fam_dic[j])
        # print(cluster)
        clusterx = {}
        for i in cluster:
            for j in range(len(i)):

                id0 = i[j]
                clusterx[id0] = j
                self.map2(axes1,((gff[id0]['end']/chro_lingth)*0.8)+0.1,(j)*0.005+0.32+(j)*0.001,0.002,family_dic[id0])
        ratio = 0.5
        for i in ks.keys():
            id1,id2 = i.split('_')[0],i.split('_')[1]
            if id1 in family and id2 in family and self.cluster0(cluster,id1,id2):
                if self.parameter == 'True':
                    ks_v = ks[i]
                    if ks_v >= float(self.ks_concern.split(',')[0]) and ks_v <= float(self.ks_concern.split(',')[1]):
                        col = return_col(ks_v,float(self.ks_concern.split(',')[0]),float(self.ks_concern.split(',')[1]))
                    else:
                        col = "#d4dcda"
                    self.plot_bez_Ks2(((gff[id1]['end']/chro_lingth)*0.8)+0.101, (clusterx[id1]) * 0.006+0.32, ((gff[id2]['end']/chro_lingth)*0.8)+0.101, (clusterx[id2])*0.006+0.32, col, ratio)
                else:
                    self.plot_bez_Ks2(((gff[id1]['end']/chro_lingth)*0.8)+0.101, (clusterx[id1]) * 0.006+0.32, ((gff[id2]['end']/chro_lingth)*0.8)+0.101, (clusterx[id2])*0.006+0.32, 'red', ratio)

    def run(self):
        lens,chr_list = read_lens(self.lens)# 染色体字典
        gff = read_gff(self.gff)# 基因字典
        family,family_dic = read_family1(self.genefamily)# 家族列表
        for i in family:# 去除不在染色体上的家族成员
            if i not in gff.keys():
                family.remove(i)
        if self.file_type == 'BLAST':# 关系列表
            blast = read_blast0(self.pair_file)
        elif self.file_type == 'WGDI':# 关系列表
            blast = read_WGDI0(self.pair_file,self.block)
        else:
            print('暂不能解析此文件格式')
            exit()
        ks = read_ks0(self.ks)# ks字典
        self.make_plot(lens,gff,family,blast,ks,family_dic)
