import ternary
import matplotlib.pyplot as pp
import numpy as np
import data_ternary
import data_udiagram
import data_ediagram

def plotTdiagram():
    scale = 100

    groups = [['grt_E11_a', 'grt_E11_b', 'grt_E11_c'], ['grt_E11_d', 'grt_E11_e', 'grt_E11_g'], ['grt_E11_j', 'grt_E11_k'],
              ['grt_Fro_b', 'grt_Fro_c', 'grt_Fro_d'], ['grt_Fro_e', 'grt_Fro_f', 'grt_Fro_g'], ['grt_Fro_n', 'grt_Fro_o', 'grt_Fro_p'], ['grt_Fro_q', 'grt_Fro_r', 'grt_Fro_s'],
              ['grt_Rane_a', 'grt_Rane_b'], ['grt_Rane_c', 'grt_Rane_d']]

    for group in groups:
        #group = [group]

        figure, tax = ternary.figure(scale=scale)

        # Draw Boundary and Gridlines
        tax.boundary(linewidth=2.0)
        tax.gridlines(color="grey", multiple=20)

        # Set Axis labels and Title
        fontsize = 12
        offset = 0.14
        #tax.set_title(group+"\n", fontsize=fontsize)
        tax.right_corner_label("Sps", fontsize=fontsize)
        tax.top_corner_label("Grs", fontsize=fontsize)
        tax.left_corner_label("Prp", fontsize=fontsize)
        #tax.left_axis_label("Left label $\\alpha^2$", fontsize=fontsize, offset=offset)
        #tax.right_axis_label("Right label $\\beta^2$", fontsize=fontsize, offset=offset)
        #tax.bottom_axis_label("Bottom label $\\Gamma - \\Omega$", fontsize=fontsize, offset=offset)

        # Draw lines parallel to the axes
        #tax.horizontal_line(16)
        #tax.left_parallel_line(10, linewidth=2., color='red', linestyle="--")
        #tax.right_parallel_line(20, linewidth=3., color='blue')

        # Draw an arbitrary line, ternary will project the points for you
        #p1 = (22, 8, 10)
        #p2 = (2, 22, 16)
        #tax.line(p1, p2, linewidth=3., marker='s', color='green', linestyle=":")

        for item in group:
            points = list(zip(*(data_ternary.data[item]['Sps'],
                                data_ternary.data[item]['Grs'],
                                data_ternary.data[item]['Prp'],
                                )))
            #print(points)
            tax.scatter(points, label=item)

            #p1 = (data_ternary.data[item][val][0] for val in ['Prp', 'Sps', 'Grs'])
            p1 = points[0]
            #p2 = (data_ternary.data[item][val][-1] for val in ['Prp', 'Sps', 'Grs'])
            p2 = points[-1]
            print(p1, p2)
            #tax.line(p1, p2, linewidth=3., marker='s', color='green', linestyle="-")
#            line.axes.annotate('',
#                               xytext=(xdata[start_ind], ydata[start_ind]),
#                               xy=(xdata[end_ind], ydata[end_ind]),
#                               arrowprops=dict(arrowstyle="->", color=color),
#                               size=size
#                               )

        tax.legend(frameon=False)

        # edit this for ax ticks
        tax.ticks(axis='lbr', multiple=20, linewidth=1, offset=0.025)

        # leave this alone
        tax.get_axes().axis('off')
        tax.clear_matplotlib_ticks()
        pp.savefig(f"ternary_{group}.png")
    tax.show()

def plotUdiagram():

    for g in data_udiagram.data.keys():

        dat = data_udiagram.data[g]

        for i, zone in enumerate(dat['zones']):
            name = f"{g}_{zone}"
            fig, ax = pp.subplots(1, num=name)
            ax.plot(dat['index'][i], dat['al4'][i], label="Al(IV)", color="blue", marker=".")
            ax.plot(dat['index'][i], dat['al6'][i], label="Al(VI)", color="orange", marker=".")
            ax.set_xlabel("No.")
            ax.set_ylabel("pfu")
            ax.set_title(f"{g}: {zone}")
            ax.legend(frameon=False)
            ax.xaxis.get_major_locator().set_params(integer=True)

            fig.savefig(f"uplot_{g}_{zone}.png")

    #pp.show()


def plotZdiagram():

    for g in data_udiagram.data.keys():

        fig, ax = pp.subplots(1, num=g)
        dat = data_udiagram.data[g]

        for i, zone in enumerate(dat['zones']):
            ax.plot(dat['al4'][i], dat['al6'][i], label=f"{zone}", marker=".")
            ax.set_xlabel("Al(IV)")
            ax.set_ylabel("Al(VI)")
            ax.set_title(g)
            ax.legend(frameon=False)
            #ax.xaxis.get_major_locator().set_params(integer=True)

        fig.savefig(f"zplot_{g}.png")

    #pp.show()

def plotEdiagram():

    #line =  [[2.75, 1], [2,.75, 0]]
    line =  [[2.75, 2.075], [1, 0]]
    labelpositions = {
            "Epi_E_a": [[2.1,0.8], [2.4,0.1]],
            "Epi_E_d": [[2.1,0.4], [2.4,0.1]],
            "Epi_E_h": [[2.1,0.5], [2.4,0.1]],
            "Epi_E_k": [[1.8,0.4], [2.4,0.1]],
            "Epi_Fro_a": [[2.1,0.5], [2.4,0.1]],
            "Epi_Fro_d": [[2.01,0.5], [2.4,0.1]],
            "Epi_Fro_g": [[2.1,0.8], [2.4,0.1]],
            "Epi_Fro_k": [[2.1,0.8], [2.4,0.1]],
            "Epi_Fro_n": [[2,0.5], [2.4,0.1]],
            "Epi_Fro_q": [[2.1,0.4], [2.4,0.1]],
            "Epi_Fro_u": [[2.08,0.3], [2.4,0.1]],
            "Epi_Rane_a": [[2.1,0.8], [2.225,0.02]],
            "Epi_Rane_d": [[1.7,0.3], [2.4,0.1]],
    }

    for dataset in data_ediagram.data:
        #print(dataset)
        firstname = dataset['labels'][0]
        fig, ax = pp.subplots(1, num=firstname)
        ax.plot(line[0], line[1], color="k")

        lens = list(map(len, dataset['al']))
        for idx in range(len(lens)):
            labelidx = sum(lens[:idx])
            label = dataset['labels'][labelidx]
            aldata = dataset['al'][idx]
            fedata = dataset['fe'][idx]

            ax.plot(aldata, fedata, linestyle="", marker="o", label=label)

        ax.annotate("Epidote", labelpositions[firstname][0])
        ax.annotate("Zoisite/Clinozoisite", labelpositions[firstname][1])
        ax.legend(frameon=False)
        ax.set_xlabel("Al")
        ax.set_ylabel("Fe+3")
        
        fig.savefig(f"eplot_{firstname}")
    pp.show()

if __name__ == "__main__":
    pass

    #plotTdiagram()
    plotUdiagram()
    plotZdiagram()
    #plotEdiagram()