#!/usr/bin/env python
import argparse
import os
import pandas as pd
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

plt.style.use('classic')
plt.rcParams.update({'legend.fancybox': True, 'legend.labelspacing': 0.25,
                     'legend.fontsize': 14, 'legend.scatterpoints': 1,
                     'font.size': 10, 'font.family': 'sans-serif',
                     'savefig.transparent': True, 'savefig.dpi': 300,
                     'savefig.bbox': 'tight',
                     'xtick.direction': 'in',
                     'figure.dpi': 80,
                     'axes.grid': True,
                     'axes.xmargin': 0, 'axes.labelsize': 16,
                     'xtick.major.size': 5, 'xtick.minor.size': 3,
                     'ps.useafm': True,
                     'pdf.use14corefonts': True})


description = """describe"""
par = argparse.ArgumentParser(description=description)
par.add_argument('xlsfile', help="input excel file")
par.add_argument('-si', '--sheet_index', default=0, type=int,
                 help="sheet index in excel file [0]")
par.add_argument('-hl', '--header', default=0, type=int,
                 help="header line number in excel file[0]")
par.add_argument('-m', '--mol_col', nargs=3, default=[0, 1, 2], type=int,
                 help="row number for mol percent [0 1 2]")
par.add_argument('-p', '--GC_col', default=3, type=int,
                 help="row number for G/C ratio [3]")
par.add_argument('-xmin', '--xmin', default=0, type=float,
                 help="x minimum [0]")
par.add_argument('-o', '--outfile', default=None,
                 help="output file[None]")
par.add_argument('-xl', '--xlabel', default="mol%",
                 help='xlabel [mol%%]')
par.add_argument('-g', '--grid', default=400, type=int,
                 help="grid size [400]")
par.add_argument('-fs', '--figsize', default=8, type=int,
                 help="figure size [8]")
par.add_argument('-ls', '--labelsize', default=16, type=float,
                 help="label font size [16]")
par.add_argument('-ts', '--ticksize', default=10, type=float,
                 help="ticks font size [10]")
par.add_argument('--labeloffset', default=6,
                 help="lebel offset [6]")
par.add_argument('--ticksoffset', default=4,
                 help="ticks offset [1]")
par.add_argument('--colormap', default='bwr',
                 help="colour color map [bwr]")
par.add_argument('--interp', default="linear",
                 choices=['linear', 'nearest', 'cubic'], help="[linear]")
par.add_argument('--nolegend', default=False,
                 action='store_true', help='not show legend [False]')
par.add_argument('--nocontour', default=False,
                 action='store_true', help='not contour plot [False]')
par.add_argument('--nocolorbar', default=False,
                 action='store_true', help='not colorbar [False]')
args = par.parse_args()

plt.rcParams.update({'font.size': args.labelsize})


def tri2car(data, k=100):
    a = data[:, 0]
    b = data[:, 1]
    c = data[:, 2]
    x = k * 0.5 * (2. * b + c) / (a + b + c)
    y = k * 0.5 * np.sqrt(3) * c / (a + b + c)
    return x, y


def tri2car_(data, k=100):
    a = data[:, 0]
    b = data[:, 1]
    c = data[:, 2]
    dx = 100 - args.xmin
    a = (a - 0)/dx * 100
    b = (b - 0)/dx * 100
    c = (c - args.xmin)/dx * 100
    x = k * 0.5 * (2. * b + c) / (a + b + c)
    y = k * 0.5 * np.sqrt(3) * c / (a + b + c)
    return x, y


def car2tri(x, y):
    if x is None or y is None:
        return None, None, None
    x, y = 0.01 * x, 0.01 * y
    c = 2 / np.sqrt(3) * y
    b = x - y / np.sqrt(3)
    a = 1 - b - c
    a = a * (100 - args.xmin)
    b = b * (100 - args.xmin)
    c = c * (100 - args.xmin) + args.xmin
    return a, b, c


def MakeFig(label, xlabel, psum):
    # 枠線の生成
    frame = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 0, 0]])
    fig, ax = plt.subplots(figsize=(args.figsize, args.figsize))
    x, y = tri2car(frame)
    ax.plot(x, y, "k-", lw=1)
    # grid線の生成
    m = 100 * np.linspace(0.1, 0.9, 9)
    g1, g2, g3 = np.zeros((9, 3)), np.zeros((9, 3)), np.zeros((9, 3))
    # test
    m = 100 * np.linspace(0, 1.0, 11)
    mx = np.linspace(args.xmin, 100, 11)
    my = np.linspace(0, 100 - args.xmin, 11)
    mz = np.linspace(0, 100 - args.xmin, 11)
    g1, g2, g3 = np.zeros((11, 3)), np.zeros((11, 3)), np.zeros((11, 3))

    g1[:, 0], g1[:, 1] = m, m[::-1]
    g2[:, 0], g2[:, 2] = m, m[::-1]
    g3[:, 1], g3[:, 2] = m, m[::-1]
    x1, y1 = tri2car(g1)
    x2, y2 = tri2car(g2)
    x3, y3 = tri2car(g3)
    x4, y4 = x1[::-1], y1[::-1]
    dashed = [2, 2]
    t1, t2 = args.ticksoffset, args.ticksoffset / np.sqrt(2)
    t1, t2 = args.ticksoffset * 0.5, args.ticksoffset * 0.5 * np.sqrt(3)
    props = {'ha': 'center', 'va': 'center', 'size': args.ticksize}
    for i, xi in enumerate(x1):
        ax.plot([x1[i], x2[i]], [y1[i], y2[i]], "k:", lw=0.2, dashes=dashed)
        ax.plot([x2[i], x3[i]], [y2[i], y3[i]], "k:", lw=0.2, dashes=dashed)
        ax.plot([x4[i], x3[i]], [y4[i], y3[i]], "k:", lw=0.2, dashes=dashed)
        v1, v2 = "%.1f" % (m[i]), "%.1f" % (m[-i - 1])
        # bottom, C
        v2 = "%.1f" % (mz[-i - 1] * 0.01 * psum)
        ax.text(x1[i] - t1, y1[i] - t2, v2, props, rotation=60)
        # left, B
        v1 = "%.1f" % (my[i] * 0.01 * psum)
        ax.text(x2[i] - t1, y2[i] + t2, v1, props, rotation=-60)
        # right, A
        v2 = "%.1f" % (mx[-i - 1] * 0.01 * psum)
        ax.text(x3[i] + args.ticksoffset, y3[i], v2, props)
    # ラベルの生成
    t1, t2 = args.labeloffset, args.labeloffset / np.sqrt(2)
    ax.text(-t2,  -t2, label[0], ha="right")
    ax.text(100 + t2, -t2, label[1], ha="left")
    ax.text(50,  50 * np.sqrt(3) + t1, label[2], ha="center", va="bottom")
    ax.text(50,  -12, xlabel, ha="center", va="bottom")
    ax.axis('off')
    offset = 1
    ax.set_xlim(-offset, 120)
    ax.set_ylim(-offset, 50 * np.sqrt(3) + offset)
    fig.gca().set_aspect('equal', adjustable='box')
    return fig, ax


def ScatterPlot(header, data, xlabel, val, psum=100, order=[1, 2, 0]):
    # cmap = plt.cm.get_cmap(args.colormap)
    cmap = mpl.colormaps.get_cmap(args.colormap)
    fig, ax = MakeFig(header[order], xlabel, psum)
    x, y = tri2car_(data[:, order])
    # contour
    if args.nocontour is False:
        xi = np.linspace(0, 100, args.grid)
        yi = np.linspace(0, 100, args.grid)
        zi = griddata((x, y), val, (xi[None, :], yi[:, None]),
                      method=args.interp)
        zi[zi == 0] = 1e-8
        alpha = 0.8
        level = np.linspace(0, 1, 20)
        plt.contourf(xi, yi, zi, level, cmap=cmap, alpha=alpha)
        # 内挿して得られる最大ポイントの表示
        # zi_ = np.nan_to_num(zi)
        # y_, x_ = divmod(np.argmax(zi_), args.grid)
        # a_, b_, c_ = car2tri(xi[x_], yi[y_])
        # s = ", ".join(header[order])
        # print("Maximum point: %s = %.2f %.2f %.2f " % (s, a_, b_, c_))

    # scatter
    bounds = np.linspace(0, 1, 11)
    norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
    # glassのプロット
    idx = (val == 1)
    if len(idx) != 0:
        ax.scatter(x[idx], y[idx], marker="o", label="glass",
                   color=cmap(0.999), edgecolor="black")
    idx = ((0 < val) & (val < 1))
    # crystalとglassの混ざったもののプロット
    if len(idx) != 0:
        ax.scatter(x[idx], y[idx], marker="^", label="glass/crystal",
                   color=cmap(0.5), edgecolor="black")
    cs = ax.scatter(x, y, c=val, cmap=cmap, norm=norm, s=0.0)
    # crystalのプロット
    idx = (val == 0)
    if len(idx) != 0:
        ax.scatter(x[idx], y[idx], marker="x", label="crystal", color=cmap(0))
    # colorbarのプロット
    if args.nocolorbar is False:
        cbar = fig.colorbar(cs, fraction=0.035, alpha=1.0)
        cbar.ax.set_ylabel('glass vol%', fontsize=14)
        cbar.ax.set_yticklabels(['%.0f' % (200 * x)
                                 for x in bounds], fontsize=12)
    if args.nolegend is False:
        ax.legend(loc="best")
    return fig


# xlsデータの読み込み
df = pd.read_excel(args.xlsfile, skiprows=args.header)
df = df.dropna()
# mol_data = mol_data/18.67*100
mol_header = df.columns
mol_data = df.iloc[:, args.mol_col].to_numpy()
val = df.iloc[:, args.GC_col].to_numpy()
print("Number of data : ", df.shape)
print("Plot header [col]: ", mol_header[args.mol_col].to_list(), args.mol_col)
print("G/C data [col]: {} [{}]".format(mol_header[args.GC_col], args.GC_col))

psum = np.sum(mol_data, axis=1)[0]
# プロット
fig_mol = ScatterPlot(mol_header, mol_data, args.xlabel, val, psum)

if args.outfile is not None:
    fig_mol.savefig(args.outfile)
    print(args.outfile, "was created.")


def onclick(event):
    a, b, c = car2tri(event.xdata, event.ydata)
    if a is None:
        return
    s = mol_header[[1, 2, 0]].to_list()
    print('button=%d %s=(%.2f %.2f %.2f)' % (event.button, s, a, b, c))


fig_mol.canvas.mpl_connect('button_press_event', onclick)
plt.show()