import os
import sys
import pandas as pd
import matplotlib.pyplot as plt
import json
from monty.json import MontyEncoder

from mp_api.client import MPRester
from emmet.core.electronic_structure import BSPathType
from pymatgen.phonon.plotter import PhononBSPlotter, PhononDosPlotter


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()

formula = None
kpath = None


def usage():
    print()
    print(f"Usage: python {sys.argv[0]} formula\n")

def update_vars():
    global formula

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

    formula = argv[1]
    if narg >= 3: kpath = argv[2]

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_material_ids(formula, mpr):
    search_results = mpr.materials.search(formula = formula, fields = [])
    if not search_results:
        print(f"No data found for {formula}")
        return None

    mid_list = []
    for res in search_results:
        mid_list.append(res.material_id)

    return mid_list


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

    mid_list = get_material_ids(formula, mpr)
    print("Material IDs:", mid_list)

    print()
    for mid in mid_list:
        print(f"Search for material ID [{mid}]...")
        ph_bs = mpr.get_phonon_bandstructure_by_material_id(mid)
        ph_dos = mpr.get_phonon_dos_by_material_id(mid)
        if ph_bs is None:
            print(f"  Phonon band data is not available: Skip")
            continue
        if ph_dos is None:
            print(f"  Phonon DOS data is not available: Skip")
            continue

        with open(f'phonon_band_{formula}_{mid}.json', 'w') as f:
            json.dump(ph_bs.as_dict(), f, cls = MontyEncoder, indent = 2)

        with open(f'phonon_dos_{formula}_{mid}.json', 'w') as f:
            json.dump(ph_dos.as_dict(), f, cls = MontyEncoder, indent = 2)

        phonon_bs_data = []
        for band in ph_bs.bands:
            for qpoint, frequency in zip(ph_bs.qpoints, band):
                phonon_bs_data.append([qpoint.frac_coords, frequency])

        phonon_bs_df = pd.DataFrame(phonon_bs_data, columns=['Q-Point', 'Frequency'])
        phonon_bs_df.to_csv(f'phonon_band_{formula}_{mid}.csv', index=False)

        phonon_dos_data = []
        for frequency, density in zip(ph_dos.frequencies, ph_dos.densities):
            phonon_dos_data.append([frequency, density])

        phonon_dos_df = pd.DataFrame(phonon_dos_data, columns=['Frequency', 'Density of States'])
        phonon_dos_df.to_csv(f'phonon_dos_{formula}_{mid}.csv', index=False)

        phonon_bs_plotter = PhononBSPlotter(ph_bs)
        plt_phonon_bs = phonon_bs_plotter.get_plot()
        fig_phonon_bs = plt_phonon_bs.figure
        fig_phonon_bs.savefig(f'phonon_band_{formula}_{mid}.png')

        phonon_dos_plotter = PhononDosPlotter()
        phonon_dos_plotter.add_dos("Total DOS", ph_dos)
        plt_phonon_dos = phonon_dos_plotter.get_plot()
        fig_phonon_dos = plt_phonon_dos.figure
        fig_phonon_dos.savefig(f'phonon_dos_{formula}_{mid}.png')

        print("JSON、CSV、PNGファイルの保存が完了しました。")


def main():
    update_vars()

    get_phonon_data()

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


if __name__ == "__main__":
    main()