"""
__author__ = "Hiroki Koiso"
__copyright__ = "Copyright 2023, Nakajima group"
__version__ = "2.1"
__maintainer__ = "Hiroki Koiso"
__email__ = "koiso.h.aa@m.titech.ac.jp"
__status__ = "Development"
__date__ = "July 22, 2023"
"""

import plotly.graph_objs as go
import plotly.offline as pyo
import pandas as pd
import numpy as np
import argparse, yaml

### Parser structure ###
desc = """
This program plots the average Grüneisen parameter at each q-point. 
Before using this program, you need to make gruneisen.yaml obtained as following:
phonopy-gruneisen orig plus minus --dim="2 2 2" --mesh="20 20 20" -c POSCAR --nomeshsym
"""
parser = argparse.ArgumentParser(description=desc)
parser.add_argument("--yaml", "-yaml", dest="name_of_yaml", default="gruneisen.yaml", type=str,
                    help="Input a gruneisen.yaml file.")
parser.add_argument("--color-bar-min", "-cmin", dest="cmin", default=-5, type=float,
                    help="Maximum value of the colorbar.")
parser.add_argument("--color-bar-max", "-cmax", dest="cmax", default=5, type=float,
                    help="Minimum value of the colorbar")
parser.add_argument("--cutoff-freq", "-cut_f", dest="cut_f", default=0.0, type=float,
                    help="the mode-Gruneisen parameters below this cutoff frequency are not used.")
parser.add_argument("--original-position", "-ori-posi", dest="original_position", action="store_true",
                    help="Do not fix q-positions.")
parser.add_argument("--fill-BZ", "-f-BZ", dest="f_BZ", action="store_true",
                    help="Fill in the first Brillouin zone.")
parser.add_argument("--relative-coordinate", "-r-c", dest="plot_rc", action="store_true",
                    help="Plot in relative coordinate.")
args = parser.parse_args()

yaml_file = args.name_of_yaml
cmin = args.cmin
cmax = args.cmax
cut_f = args.cut_f
fix_position = not args.original_position
f_BZ = args.f_BZ
plot_rc = args.plot_rc


### The below function was made by Qijing Zheng. ###
### http://staff.ustc.edu.cn/~zqj/posts/howto-plot-brillouin-zone/ ###
def get_brillouin_zone_3d(cell):
    """ Generate the Brillouin Zone of a given cell.
    The BZ is the Wigner-Seitz cell of the reciprocal lattice, which can be constructed by Voronoi decompositionto the reciprocal lattice. 
    A Voronoi diagram is a subdivision of the space into the nearest neighborhoods of a given set of points. 

    https://en.wikipedia.org/wiki/Wigner%E2%80%93Seitz_cell
    https://docs.scipy.org/doc/scipy/reference/tutorial/spatial.html#voronoi-diagrams

    Parameters
    ----------
    cell : ndarray, shape=(3, 3), raw vectors [b1, b2, b3]^T
        reciprocal lattice vector  

    Returns
    ----------
    vertices : tuple
        vertices of BZ
    ridges : list of list
        ridges of BZ
    faces : list of list
        faces of BZ

    """
    cell = np.asarray(cell, dtype=float)
    assert cell.shape == (3, 3)

    px, py, pz = np.tensordot(cell, np.mgrid[-1:2, -1:2, -1:2], axes=[0, 0])
    points = np.c_[px.ravel(), py.ravel(), pz.ravel()]
    from scipy.spatial import Voronoi
    vor = Voronoi(points)

    bz_facets = []
    bz_ridges = []
    bz_vertices = []

    for pid, rid in zip(vor.ridge_points, vor.ridge_vertices):
        # WHY 13 ????
        # The Voronoi ridges/facets are perpendicular to the lines drawn between the
        # input points. The 14th input point is [0, 0, 0].
        if(pid[0] == 13 or pid[1] == 13):
            bz_ridges.append(vor.vertices[np.r_[rid, [rid[0]]]])
            bz_facets.append(vor.vertices[rid])
            bz_vertices += rid

    bz_vertices = list(set(bz_vertices))

    return vor.vertices[bz_vertices], bz_ridges, bz_facets


def q_position_fixer(q_position, fix_position):
    """ Fix q-points to be within half of the first Brillouin zone. """
    if fix_position == True:
        if q_position[2] <= -0.5:
            q_position = q_position * (-1)
    else:
        pass
    return q_position


def get_average_gruneisen_Brillouin_zone(yaml_file, cut_f, fix_position):
    """ Calculate average Gruneisen parameter at each q-point.

    Parameters
    ----------
    yaml_file : str
        file name of gruneisen.yaml
    cut_f : float
        cut-off frequency
    fix_position : bool
        whether fix q-points to be within half of the first Brillouin zone or do not.

    Returns
    ----------
    list of plot data : ndarray
        shape=((num of q-positions), 7)
    reciprocal_lattice : ndarray, shape=(3, 3)
        raw vectors [b1, b2, b3]^T

    """
    with open(yaml_file, 'r') as file:
        config = yaml.safe_load(file)
    nqpoint = config['nqpoint']
    natom = config['natom']
    nband = 3 * natom
    rec_cell = np.array(config['reciprocal_lattice'])

    data_array = []
    for q in range(nqpoint):
        q_position = q_position_fixer(np.array(config['phonon'][q]['q-position']), fix_position)
        q_position_ac = np.dot(q_position, rec_cell)
        gamma_list =[]
        for band in range(nband):
            gruneisen = config['phonon'][q]['band'][band]['gruneisen']
            frequency = config['phonon'][q]['band'][band]['frequency']
            ### Remove phonon mode below cutoff frequency. ###
            if frequency > cut_f:
                gamma_list.append(gruneisen)
            else:
                pass
        gruneisen_average = np.mean(gamma_list)
        data = [q_position[0], q_position[1], q_position[2], q_position_ac[0], q_position_ac[1], q_position_ac[2],  gruneisen_average]
        data_array.append(data)
    return np.array(data_array), rec_cell


data_array, rec_cell = get_average_gruneisen_Brillouin_zone(yaml_file, cut_f, fix_position)
if f_BZ:
    add_data = data_array.copy()
    add_data[:, list(range(0, 6))] *= (-1)
    data_array = np.concatenate([data_array, add_data])

""" Save data """
with open('data_average_gruneisen_BZ.csv', "w") as fp:
    fp.write('qx_rc,qy_rc,qz_rc,qx_ac,qy_ac,qz_ac,gruneisen_average\n')
    np.savetxt(fp, data_array, fmt='%f', delimiter=',')



""" plot """
df_plot = pd.read_csv('data_average_gruneisen_BZ.csv')

def axis_setting(label, plot_rc):
    axis_setting = dict(
                    showticklabels=plot_rc,
                    title=dict(text=label,
                    font=dict(family='Arial',size=30,color='black')),
                    tickfont=dict(family='Arial',size=16,color='black')
    )
    return axis_setting

if plot_rc:
    qx=df_plot['qx_rc']
    qy=df_plot['qy_rc']
    qz=df_plot['qz_rc']
else: 
    qx=df_plot['qx_ac']
    qy=df_plot['qy_ac']
    qz=df_plot['qz_ac']

colorbar_layout=dict(title=dict(text='&#947;<sub>q</sub>',
                                 font=dict(family='Times New Roman',size=35,color='black')),
                                 tickfont=dict(family='Arial',size=24,color='black'),
                                 len=0.9,
                                 ticks="outside",
                                 tickwidth=2
)
marker_setting = dict(
        size=5,
        color=df_plot['gruneisen_average'],
        #colorscale='Picnic',
        colorscale=[[0.0, 'rgb(0, 0, 255)'], [(0-cmin)/(cmax-cmin), 'rgb(255, 255, 255)'], [1.0, 'rgb(255, 0, 0)']],
        cmin=cmin,
        cmax=cmax,
        opacity=1.0,
        colorbar=colorbar_layout
)
trace = go.Scatter3d(
    x=qx,
    y=qy,
    z=qz,
    customdata=np.stack((df_plot['qx_rc'],df_plot['qy_rc'],df_plot['qz_rc'],df_plot['gruneisen_average']), axis=-1),
    mode='markers',
    hovertemplate='q-position: (%{customdata[0]}, %{customdata[1]}, %{customdata[2]})<br>&#947;<sub>q</sub> average: %{customdata[3]}',
    marker=marker_setting
)

layout = go.Layout(
    margin=dict(l=0, r=0, b=0, t=0),
    showlegend=False,
    scene=dict(
        xaxis=axis_setting('<i>k<sub>x</sub></i>', plot_rc),
        yaxis=axis_setting('<i>k<sub>y</sub></i>', plot_rc),
        zaxis=axis_setting('<i>k<sub>z</sub></i>', plot_rc)
    )
)

fig = go.Figure(data=trace, layout=layout)

if not plot_rc:
    v, e, f = get_brillouin_zone_3d(rec_cell)
    
    # Base
    basis_clrs = ['red', 'green', 'blue']
    basis_labs = ['<i>b<sub>1</sub></i>', '<i>b<sub>2</sub></i>', '<i>b<sub>3</sub></i>']
    for ii, basis in enumerate(rec_cell):
        bx, by, bz = basis
        fig.add_trace(
                go.Scatter3d(
                    x=[0, bx],
                    y=[0, by],
                    z=[0, bz],
                    opacity=0.8,
                    hoverinfo='skip',
                    mode='lines+text',
                    line=dict(
                        color=basis_clrs[ii],
                        width=6,
                    ),
                    text=['', basis_labs[ii]],
                    textfont=dict(color=basis_clrs[ii], size=30),
                )
            )
    # Edige
    for edge in e:
        fig.add_trace(
                go.Scatter3d(
                    x=edge[:, 0],
                    y=edge[:, 1],
                    z=edge[:, 2],
                    opacity=0.8,
                    hoverinfo='skip',
                    mode='lines',
                    line=dict(
                        color='black',
                        width=5,
                    ),
                )
            )
#    # Vartices
#    v_rc = np.dot(v, np.linalg.inv(rec_cell))
#    fig.add_trace(
#            go.Scatter3d(
#                x=v[:, 0],
#                y=v[:, 1],
#                z=v[:, 2],
#                customdata=v_rc,
#                hovertemplate='q-position: (%{customdata[0]:.2f}, %{customdata[1]:.2f}, %{customdata[2]:.2f})',
#                opacity=1,
#                mode='markers',
#                marker=dict(
#                    color='black',
#                    size=3,
#                ),
#            )
#        )

fig.update_layout(
    showlegend=False,
    scene=dict(
        xaxis_visible=False,
        yaxis_visible=False,
        zaxis_visible=False),
    )


pyo.plot(fig, filename=f'average_gruneisen_BZ.html')
