#!/usr/bin/env python
# -*- coding: UTF-8 -*-

"""
Functions in this script produce figures in various manuscripts.
"""

import os.path as op
import sys
import logging

import numpy as np

from jcvi.graphics.base import Polygon, normalize_axes, panel_labels, plt, savefig
from jcvi.graphics.glyph import GeneGlyph, RoundRect, TextCircle, DoubleSquare, plot_cap
from jcvi.graphics.karyotype import Karyotype
from jcvi.graphics.synteny import Synteny, draw_gene_legend
from jcvi.apps.base import OptionParser, ActionDispatcher, fname


def main():

    actions = (
        # Epoch paper (Woodhouse et al., 2012 Plant Cell)
        ("epoch", "show the methods used in epoch paper"),
        # Cotton paper (Paterson et al., 2012 Nature)
        ("cotton", "plot cotton macro- and micro-synteny (requires data)"),
        # Amborella paper (Albert et al., 2013 Science)
        ("amborella", "plot amborella macro- and micro-synteny (requires data)"),
        # Mt4.0 paper (Tang et al., 2014 BMC Genomics)
        ("mtdotplots", "plot Mt3.5 and Mt4.0 side-by-side"),
        # Oropetium paper (Vanburen et al., 2015 Nature)
        ("oropetium", "plot oropetium micro-synteny (requires data)"),
        # Pomegranate paper (Qin et al., 2017 Plant Journal)
        ("pomegranate", "plot pomegranate macro- and micro-synteny (requires data)"),
        # Unpublished
        ("birch", "plot birch macro-synteny (requires data)"),
        ("litchi", "plot litchi micro-synteny (requires data)"),
        ("utricularia", "plot utricularia micro-synteny (requires data)"),
        (
            "waterlilyGOM",
            "waterlily phylogeny and related infographics (requires data)",
        ),
    )
    p = ActionDispatcher(actions)
    p.dispatch(globals())


def waterlilyGOM(args):
    """
    %prog mcmctree.tre table.csv

    Customized figure to plot phylogeny and related infographics.
    """
    from jcvi.graphics.tree import (
        LeafInfoFile,
        WGDInfoFile,
        draw_tree,
        parse_tree,
        draw_wgd_xy,
    )
    from jcvi.graphics.table import CsvTable, draw_table

    p = OptionParser(waterlilyGOM.__doc__)
    opts, args, iopts = p.set_image_options(args, figsize="12x9")

    if len(args) != 2:
        sys.exit(not p.print_help())

    (datafile, csvfile) = args
    outgroup = ["ginkgo"]

    logging.debug("Load tree file `{0}`".format(datafile))
    t, hpd = parse_tree(datafile)

    pf = datafile.rsplit(".", 1)[0]

    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])

    margin, rmargin = 0.15, 0.19  # Left and right margin
    leafinfo = LeafInfoFile("leafinfo.csv").cache
    wgdinfo = WGDInfoFile("wgdinfo.csv").cache
    groups = "Monocots,Eudicots,ANA-grade,Gymnosperms"

    draw_tree(
        root,
        t,
        hpd=hpd,
        margin=margin,
        rmargin=rmargin,
        supportcolor=None,
        internal=False,
        outgroup=outgroup,
        leafinfo=leafinfo,
       
 wgdinfo=wgdinfo,
        geoscale=True,
        groups=groups.split(","),
    )

    # Bottom right show legends for the WGD circles
    pad = 0.02
    ypad = 0.04
    xstart = 1 - rmargin + pad
    ystart = 0.2
    waterlily_wgdline = wgdinfo["waterlily"][0]
    ypos = ystart - 2 * ypad
    draw_wgd_xy(root, xstart
, ypos, waterlily_wgdline)
    root.text(
        xstart + pad,
        ypos,
        "Nymphaealean WGD",
        color=waterlily_wgdline.color,
        va="center",
    )
    other_wgdline = wgdinfo["banana"][0]
    ypos = ystart - 3 * ypad
    draw_wgd_xy(root, xstart, ypos, other_wgdline)
    root.text(
        xstart + pad,
        ypos,
        "Other known WGDs",
        color=other_wgdline.color,
        va="center",
    )

    # Top left draw the comparison table
    csv_table = CsvTable(csvfile)
    draw_table(
        root,
        csv_table,
        extent=(0.02, 0.44, 0.55, 0.985),
        stripe_color="lavender",
        yinflation=iopts.w / iopts.h,
    )

    normalize_axes(root)

    image_name = pf + "." + iopts.format
    savefig(image_name, dpi=iopts.dpi, iopts=iopts)


def pomegranate(args):
    """
    %prog cotton seqids karyotype.layout mcscan.out all.bed synteny.layout

    Build a figure that calls graphics.karyotype to illustrate the high ploidy
    of WGD history of pineapple genome. The script calls both graphics.karyotype
    and graphic.synteny.
    """
    p = OptionParser(pomegranate.__doc__)
    opts, args, iopts = p.set_image_options(args, figsize="9x7")

    if len(args) != 5:
        sys.exit(not p.print_help())

    seqidsfile, klayout, datafile, bedfile, slayout = args

    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])

    Karyotype(fig, root, seqidsfile, klayout)
    Synteny(fig, root, datafile, bedfile, slayout)

    # legend showing the orientation of the genes
    draw_gene_legend(root, 0.42, 0.52, 0.48)

    labels = ((0.04, 0.96, "A"), (0.04, 0.52, "B"))
    panel_labels(root, labels)

    root.set_xlim(0, 1)
    root.set_ylim(0, 1)
    root.set_axis_off()

    pf = "pomegranate-karyotype"
    image_name = pf + "." + iopts.format
    savefig(image_name, dpi=iopts.dpi, iopts=iopts)


def utricularia(args):
    from jcvi.graphics.synteny import main as synteny_main

    p = OptionParser(synteny_main.__doc__)
    p.add_option("--switch", help="Rename the seqid with two-column file")
    opts, args, iopts = p.set_image_options(args, figsize="8x7")

    if len(args) != 3:
        sys.exit(not p.print_help())

    datafile, bedfile, layoutfile = args
    switch = opts.switch

    pf = datafile.rsplit(".", 1)[0]
    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])

    s = Synteny(
        fig, root, datafile, bedfile, layoutfile, loc_label=False, switch=switch
    )
    light = "lightslategrey"
    RoundRect(root, (0.02, 0.69), 0.96, 0.24, fill=False, lw=2, ec=light)
    RoundRect(root, (0.02, 0.09), 0.96, 0.48, fill=False, lw=2, ec=light)
    za, zb = s.layout[1].ratio, s.layout[-1].ratio  # zoom level
    if za != 1:
        root.text(
            0.96,
            0.89,
            "{}x zoom".format(za).replace(".0x", "x"),
            color=light,
            ha="right",
            va="center",
            size=14,
        )
    if zb != 1:
        root.text(
            0.96,
            0.12,
            "{}x zoom".format(zb).replace(".0x", "x"),
            color=light,
            ha="right",
            va="center",
            size=14,
        )

    # legend showing the orientation of the genes
    draw_gene_legend(root, 0.22, 0.3, 0.64, text=True)

    root.set_xlim(0, 1)
    root.set_ylim(0, 1)
    root.set_axis_off()

    image_name = pf + "." + iopts.format
    savefig(image_name, dpi=iopts.dpi, iopts=iopts)


def join_nodes(
    root, coords, a, b, x, slope=2.4, fc="lightslategray", rectangle=True, circle=True
):
    # Join node a and b to make an internal node
    ax, ay = coords[a]
    bx, by = coords[b]
    if ay < by:
        ax, ay, bx, by = bx, by, ax, ay
    if rectangle:
        nx, ny = x, (ay + by) / 2
        root.plot((nx, ax), (ay, ay), lw=2, color=fc)
        root.plot((nx, bx), (by, by), lw=2, color=fc)
        root.plot((nx, nx), (ay, by), lw=2, color=fc)
    else:
        dx = (abs(ay - by) / slope - abs(ax - bx)) / 2
        nx = max(ax, bx) + dx
        ny = by + (nx - bx) * slope
        root.plot((nx, ax), (ny, ay), lw=2, color=fc)
        root.plot((nx, bx), (ny, by), lw=2, color=fc)
    if circle:
        DoubleSquare(root, nx, ny, fc=fc)
    return nx, ny


def branch_length(ax, start, end, text, ha="left", va="bottom", color="r"):
    xs, ys = start
    xe, ye = end
    text = r"$\mathsf{" + text + "}$"
    ax.text((xs + xe) / 2, (ys + ye) / 2, text, ha=ha, va=va, color=color)


def birch(args):
    """
    %prog birch seqids layout

    Plot birch macro-synteny, with an embedded phylogenetic tree to the right.
    """
    p = OptionParser(birch.__doc__)
    opts, args, iopts = p.set_image_options(args, figsize="8x6")

    if len(args) != 2:
        sys.exit(not p.print_help())

    seqids, layout = args
    fig = plt.figure(1, (iopts.w, iopts.h))
    root = f