import numpy as np
from pymatgen.io.vaspio.vasp_output import Vasprun
from pymatgen.electronic_structure.core import Spin

import plotly.plotly as pltly      # plotting functions
import plotly.tools as tls         # plotly tools
from plotly.graph_objs import Line, Layout, Scatter, Annotation, Annotations, XAxis, YAxis, Legend, Marker

dosxml = "/home/.../DOS/vasprun.xml"
bxml = "/home/.../Bands/vasprun.xml"
kpoints = "/home/.../Bands/KPOINTS"

dosrun = Vasprun(dosxml)
Mn3d_dos = dosrun.complete_dos.get_element_spd_dos("Mn")["D"]
O2s_dos = dosrun.complete_dos.get_element_spd_dos("O")["S"]
O2p_dos = dosrun.complete_dos.get_element_spd_dos("O")["P"]

run = Vasprun(bxml, parse_projected_eigen=True)
bands = run.get_band_structure(kpoints, line_mode=True, efermi=dosrun.efermi)
atom_select = {"Mn": ["d"], "O": ["s", "p"]}
el_spd_bands = bands.get_projections_on_elts_and_orbitals(atom_select)

emin = 1e100
emax = -1e100
for spin in bands.bands:
    for b in range(bands.nb_bands):
        emin = min(emin, min(bands.bands[spin][b]))
        emax = max(emax, max(bands.bands[spin][b]))

emin -= bands.efermi + 1
#emax -= bands.efermi - 1
emax = 10
print(emin, emax)

dosbandfig = tls.make_subplots(rows=1, cols=3)
bandsubplot = 2 # band diagram in the middle

for spin in [Spin.up, Spin.down]:
    if spin == Spin.up:
        subplot = 1
        showlegend = True
    else:
        subplot = 3
        showlegend = False

    # Density of States
    # -----------------

    # total dos
    trace_tdos_up = Scatter(
        x=dosrun.tdos.densities[spin],
        y=dosrun.tdos.energies - dosrun.efermi,
        mode="lines",
        name="total DOS",
        line=Line(color="#444444", width=1),
        fill="tozeroy",
        showlegend=showlegend
    )
    dosbandfig.append_trace(trace_tdos_up, 1, subplot)

    # O 2s contribution to the total DOS
    trace_O2s = Scatter(
        x=O2s_dos.densities[spin],
        y=dosrun.tdos.energies - dosrun.efermi,
        mode="lines",
        name="O 2s",
        line=Line(color="red", width=1),
        showlegend=showlegend
    )
    dosbandfig.append_trace(trace_O2s, 1, subplot)

    # O 2p contribution to the total DOS
    trace_O2p = Scatter(
        x=O2p_dos.densities[spin],
        y=dosrun.tdos.energies - dosrun.efermi,
        mode="lines",
        name="O 2p",
        line=Line(color="#00FF00", width=1),
        showlegend=showlegend
    )
    dosbandfig.append_trace(trace_O2p, 1, subplot)

    # Mn 3d contribution to the total DOS
    trace_Mn3s = Scatter(
        x=Mn3d_dos.densities[spin],
        y=dosrun.tdos.energies - dosrun.efermi,
        mode="lines",
        name="Mn 3d",
        line=Line(color="blue", width=1),
        showlegend=showlegend
    )
    dosbandfig.append_trace(trace_Mn3s, 1, subplot)

    # Diagramme de Bandes
    # -------------------
    contrib = np.zeros((bands.nb_bands, len(bands.kpoints), 3))
    for b in range(bands.nb_bands):
        for k in range(len(bands.kpoints)):
            tot = 0
            for oa in el_spd_bands[spin][b][k].values():
                tot += sum([coef**2 for coef in oa.values()])
            
            if tot != 0.0:
                contrib[b, k, 0] = el_spd_bands[spin][b][k]["O"]["s"]**2 / tot
                contrib[b, k, 1] = el_spd_bands[spin][b][k]["O"]["p"]**2 / tot
                contrib[b, k, 2] = el_spd_bands[spin][b][k]["Mn"]["d"]**2 / tot

    for b in range(bands.nb_bands):
        eband = [e - bands.efermi for e in bands.bands[spin][b]]
        for k in range(len(bands.kpoints) - 1):
            red, green, blue = [
                int(255 * (contrib[b, k, i] + contrib[b, k+1, i])/2) for i in range(3)]
            line_trace = Scatter(
                    x=[k, k+1],
                    y=[eband[k], eband[k+1]],
                    mode="lines",
                    showlegend=False
            )
            if spin == Spin.up:
                line_trace["mode"] = "lines"
                line_trace["line"] = Line(
                    color="rgb({}, {}, {})".format(red, green, blue), 
                    width=1)
            else:
                line_trace["mode"] = "markers"
                line_trace["marker"] = Marker(
                    color="rgb({}, {}, {})".format(red, green, blue), 
                    symbol="dot", size=0.7)
            dosbandfig.append_trace(line_trace, 1, bandsubplot)

pltly.image.ishow(dosbandfig, width=975, height=650)

labels = [ r"$L$", r"$\Gamma$", r"$X$", r"$U,K$", r"$\Gamma$" ]

step = len(bands.kpoints) / (len(labels) - 1)
for i, label in enumerate(labels):
    vline = Scatter(
                x=[i * step, i * step],
                y=[emin, emax],
                mode="lines",
                line=Line(color="#111111", width=1),
                showlegend=False
            )
    dosbandfig.append_trace(vline, 1, bandsubplot)

annotations = list()
for i, label in enumerate(labels):
    annotations.append(
        Annotation(
            x=i * step, y=emin,                                                                   
            xref="x2", yref="y2",
            text=label,
            xanchor="center", yanchor="top",
            showarrow=False
        )
    )

annotations.append(
    Annotation(
        x=len(bands.kpoints) / 2, y=emax,
        xref="x2", yref="y1",
        text="Bands Diagram",
        xanchor="center", yanchor="bottom",
        showarrow=False
    )
)
annotations.append(
    Annotation(
        x=2.25, y=emax,
        xref="x1", yref="y1",
        text="DOS up",
        xanchor="center", yanchor="bottom",
        showarrow=False
    )
)
annotations.append(
    Annotation(
        x=2.25, y=emax,
        xref="x3", yref="y1",
        text="DOS down",
        xanchor="center", yanchor="bottom",
        showarrow=False
    )
)

dosbandfig["layout"].update(
    Layout(
        title="Bands diagram and density of states of MnO",
        legend=Legend(
            x=1., y=.45,
            xanchor="right", yanchor="top",
            bordercolor='#333', borderwidth=1
        ),  
        xaxis1=XAxis( 
            domain=[.0, .25],
            ticks="inside",
            showticklabels=False,
            range=[4.1, 0.01],
        ),      
        yaxis1=YAxis(
            title="$E - E_f \quad / \quad \\text{eV}$",
            ticks="inside", 
            range=[emin, emax],
        ),  
        xaxis2=XAxis(
            domain=[.250, .75],
            range=[0, len(bands.kpoints)],
            showticklabels=False,
            ticks="", 
        ),      
        yaxis2=YAxis(
            range=[emin, emax],
            ticks="inside",
            showticklabels=False
        ),  
        xaxis3=XAxis(
            domain=[.752, 1.],
            ticks="inside",
            showticklabels=False,
            range=[.01, 4.1],
        ),
        yaxis3=YAxis(
            side="right",
            title="$E - E_f \quad / \quad \\text{eV}$",
            ticks="inside",
            range=[emin, emax]
        ),
        annotations=Annotations(annotations)
    )
)

axisdefault = {"showgrid": True,                                                                  
               "showline": True,
               "mirror": "allticks",
               "linewidth": 1,
               "tickwidth": 1}
for obj in ["yaxis1", "yaxis2", "yaxis3", "xaxis1", "xaxis2", "xaxis3"]:
    dosbandfig["layout"][obj].update(axisdefault)

#pltly.image.ishow(dosbandfig, width=975, height=650)
plot_url = pltly.plot(dosbandfig, filename="MnO_DOS_bands", auto_open=False)
tls.embed(plot_url)