import json
import numpy as np
import pandas as pd
from typing import List, Tuple, Dict, Optional


def coord2df(coord: List[Tuple[str, List]]) -> pd.DataFrame:
    if isinstance(coord, pd.DataFrame):
        return coord
    else:
        return pd.DataFrame(coord, columns=["element", "coord_x", "coord_y", "coord_z"])


def dict2jsonfile(dict_obj: Dict, filepath: str) -> None:
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(dict_obj, f, ensure_ascii=False, indent=2)


def dict2dataframe(
    data_dict: Dict[str, Dict], smiles_list: Optional[List[str]] = None, name_list: Optional[List[str]] = None
) -> pd.DataFrame:
    if not data_dict:
        return pd.DataFrame()
    df = pd.DataFrame.from_dict(data_dict, orient="index")
    if name_list:
        df.index = name_list[: len(df)]
        df.index.name = "tag"

    if smiles_list:
        df.insert(0, "SMILES", smiles_list[: len(df)])

    return df


def df2xyz(df: pd.DataFrame, filepath: str, tag: str) -> None:
    required_columns = {"element", "coord_x", "coord_y", "coord_z"}
    assert required_columns.issubset(df.columns), "DataFrame must contain element and coordinate columns."

    with open(filepath, "w") as f:
        f.write(f"{len(df)}\n")
        f.write(f"file {tag} which generated by qspoc.\n")
        for _, row in df.iterrows():
            f.write(f"{row['element']} {row['coord_x']} {row['coord_y']} {row['coord_z']}\n")


def xyz2gjf(
    xyz_filepath: str,
    gjf_filepath: str,
    tag: str,
    calc_command: str,
    charge: int = 0,
    multiplicity: int = 1,
    nprocs: int = 4,
    memory: str = "2GB",
    chk_path: str = None,
) -> None:
    with open(xyz_filepath, "r") as f:
        lines = f.readlines()

    geometry = [line.strip() for line in lines[2:] if line.strip()]

    with open(gjf_filepath, "w") as f:
        if chk_path:
            f.write(f"%chk={chk_path}\n")
        f.write(f"%mem={memory}\n")
        f.write(f"%nprocshared={nprocs}\n")
        f.write(f"#p {calc_command}\n\n")
        f.write(f"{tag} generated by qspoc.\n\n")
        f.write(f"{charge} {multiplicity}\n")
        for line in geometry:
            f.write(f"{line}\n")
        f.write("\n")


def xyz2coord(xyz_filepath: str) -> pd.DataFrame:
    with open(xyz_filepath, "r") as f:
        lines = f.readlines()
    assert len(lines) > 2, "Invalid XYZ file format."
    coord = [[line.split()[0], np.array(line.split()[1:], dtype=float)] for line in lines[2:]]

    return pd.DataFrame(coord, columns=["element", "coord"])


def dataframe2csv(df: pd.DataFrame, filepath: str, **kwargs) -> None:
    df.to_csv(filepath, **kwargs)


def dataframe2xlsx(df: pd.DataFrame, filepath: str, **kwargs) -> None:
    df.to_excel(filepath, **kwargs)


def csv2dataframe(filepath: str, **kwargs) -> pd.DataFrame:
    return pd.read_csv(filepath, **kwargs)


def xlsx2dataframe(filepath: str, **kwargs) -> pd.DataFrame:
    return pd.read_excel(filepath, **kwargs)
