import os
import sys
import pandas as pd
import matplotlib.pyplot as plt
import json

from mp_api.client import MPRester
from emmet.core.thermo import ThermoType
from pymatgen.analysis.phase_diagram import PDPlotter
from monty.json import MontyEncoder


API_KEY = os.getenv('MP_APIKEY')
if API_KEY is None:
        print("\nError: Can not get MP API Key from the environment var MP_APIKEY\n")
        exit()


chemsys = None  # "Li-Fe-O"


def usage():
    print()
    print(f"Usage: python {sys.argv[0]} chemsys (Li-Fe-O etc)\n")

def update_vars():
    global chemsys

    argv = sys.argv
    narg = len(argv)
    if narg <= 1:
        print("\nError: Chemical system must be given as the first arg\n")
        input("Press ENTER to terminate>>\n")
        exit()

    chemsys = argv[1]

def get_MPReser(API_KEY = None):
    if not API_KEY: API_KEY = os.getenv('MP_APIKEY')
    if API_KEY is None:
        print("\nError: Can not get MP API Key from the environment var MP_APIKEY\n")
        return None

    mpr = MPRester(API_KEY)
    if mpr is None:
        print(f"\nError: Can not get MPRester using the given API_KEY [{API_KEY}]\n")
        return None

    return mpr


def get_phase_diagram_data():
    mpr = get_MPReser()
    if not mpr: exit()

    print()
    pd_obj = mpr.materials.thermo.get_phase_diagram_from_chemsys(chemsys = chemsys, 
                                                   thermo_type=ThermoType.GGA_GGA_U_R2SCAN)
#    pd_obj = mpr.materials.thermo.get_phase_diagram_from_chemsys(chemsys="Li-Fe-O", 
#                                                   thermo_type=ThermoType.GGA_GGA_U)
#    pd_obj = mpr.materials.thermo.get_phase_diagram_from_chemsys(chemsys="Li-Fe-O", 
#                                                   thermo_type=ThermoType.R2SCAN)

    if pd_obj is None:
        print(f"  PHase diagram data is not available: Skip")
        return None

#    print("pd_obj=", pd_obj)

    with open(f'phase_diagram_{chemsys}.json', 'w') as f:
        json.dump(pd_obj.as_dict(), f, cls=MontyEncoder, indent = 2)

    pd_data = []
    for entry in pd_obj.all_entries:
        pd_data.append([entry.composition, entry.energy])

    outfile = f'phase_diagram_{chemsys}.csv'
    print(f"Save CSV file to [{outfile}]")
    pd_df = pd.DataFrame(pd_data, columns=['Composition', 'Energy'])
    pd_df.to_csv(outfile, index = False)

    pd_plotter = PDPlotter(pd_obj)
    plt_pd = pd_plotter.get_plot()
    plt_pd.show()

    outfile = f'phase_diagram_{chemsys}.html'
    if 'plotly' in str(type(plt_pd)):
#        print("Plotly is used:")
#        plt_pd.write_image(outfile)
        print(f"Save HTML file to [{outfile}]")
        plt_pd.write_html(outfile)


def main():
    update_vars()

    get_phase_diagram_data()

    usage()
    print()
    input("Press ENTER to terminate>>\n")
    exit()


if __name__ == "__main__":
    main()