"""
This module contains the functions for reading the magephem data from the txt and h5 files from RBSP ECT website.
https://rbsp-ect.newmexicoconsortium.org/data_pub/
@Author: Xiangning Chu, 2025-04-07
"""

import linecache

import h5py
import numpy as np
import pandas as pd
from pytplot import store_data

name_332 = [
    "Time",
    "Date",
    "DOY",
    "UTC",
    "JulianDate",
    "GpsTime",
    "DipoleTiltAngle",
    "InOut",
    "Orbnumber",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "Rgeod_Lat",
    "Rgeod_Lon",
    "Rgeod_Height",
    "Rgsm_x",
    "Rgsm_y",
    "Rgsm_z",
    "Rsm_x",
    "Rsm_y",
    "Rsm_z",
    "Rgei_x",
    "Rgei_y",
    "Rgei_z",
    "Rgse_x",
    "Rgse_y",
    "Rgse_z",
    "CDMAG_MLAT",
    "CDMAG_MLON",
    "CDMAG_MLT",
    "CDMAG_R",
    "EDMAG_MLAT",
    "EDMAG_MLON",
    "EDMAG_MLT",
    "EDMAG_R",
    "IntModel",
    "ExtModel",
    "Kp",
    "Dst",
    "Bsc_gsm_x",
    "Bsc_gsm_y",
    "Bsc_gsm_z",
    "Bsc_gsm_mag",
    "FieldLineType",
    "S_sc_to_pfn",
    "S_sc_to_pfs",
    "S_pfs_to_Bmin",
    "S_Bmin_to_sc",
    "S_total",
    "d2B_ds2",
    "Sb0",
    "RadiusOfCurv",
    "Pfn_gsm_x",
    "Pfn_gsm_y",
    "Pfn_gsm_z",
    "Pfn_geo_x",
    "Pfn_geo_y",
    "Pfn_geo_z",
    "Pfn_geod_Lat",
    "Pfn_geod_Lon",
    "Pfn_geod_Height",
    "Pfn_CD_MLAT",
    "Pfn_CD_MLON",
    "Pfn_CD_MLT",
    "Pfn_ED_MLAT",
    "Pfn_ED_MLON",
    "Pfn_ED_MLT",
    "Bfn_geo_x",
    "Bfn_geo_y",
    "Bfn_geo_z",
    "Bfn_geo_mag",
    "Bfn_gsm_x",
    "Bfn_gsm_y",
    "Bfn_gsm_z",
    "Bfn_gsm_mag",
    "Loss_Cone_Alpha_n",
    "Pfs_gsm_x",
    "Pfs_gsm_y",
    "Pfs_gsm_z",
    "Pfs_geo_x",
    "Pfs_geo_y",
    "Pfs_geo_z",
    "Pfs_geod_Lat",
    "Pfs_geod_Lon",
    "Pfs_geod_Height",
    "Pfs_CD_MLAT",
    "Pfs_CD_MLON",
    "Pfs_CD_MLT",
    "Pfs_ED_MLAT",
    "Pfs_ED_MLON",
    "Pfs_ED_MLT",
    "Bfs_geo_x",
    "Bfs_geo_y",
    "Bfs_geo_z",
    "Bfs_geo_mag",
    "Bfs_gsm_x",
    "Bfs_gsm_y",
    "Bfs_gsm_z",
    "Bfs_gsm_mag",
    "Loss_Cone_Alpha_s",
    "Pmin_gsm_x",
    "Pmin_gsm_y",
    "Pmin_gsm_z",
    "Bmin_gsm_x",
    "Bmin_gsm_y",
    "Bmin_gsm_z",
    "Bmin_gsm_mag",
    "Lsimple",
    "InvLat",
    "Lm_eq",
    "InvLat_eq",
    "BoverBeq",
    "MlatFromBoverBeq",
    "M_used",
    "M_ref",
    "M_igrf",
    "Hmin_90",
    "Hmin_85",
    "Hmin_80",
    "Hmin_75",
    "Hmin_70",
    "Hmin_65",
    "Hmin_60",
    "Hmin_55",
    "Hmin_50",
    "Hmin_45",
    "Hmin_40",
    "Hmin_35",
    "Hmin_30",
    "Hmin_25",
    "Hmin_20",
    "Hmin_15",
    "Hmin_10",
    "Hmin_5",
    "Hmin_GeodLat_90",
    "Hmin_GeodLat_85",
    "Hmin_GeodLat_80",
    "Hmin_GeodLat_75",
    "Hmin_GeodLat_70",
    "Hmin_GeodLat_65",
    "Hmin_GeodLat_60",
    "Hmin_GeodLat_55",
    "Hmin_GeodLat_50",
    "Hmin_GeodLat_45",
    "Hmin_GeodLat_40",
    "Hmin_GeodLat_35",
    "Hmin_GeodLat_30",
    "Hmin_GeodLat_25",
    "Hmin_GeodLat_20",
    "Hmin_GeodLat_15",
    "Hmin_GeodLat_10",
    "Hmin_GeodLat_5",
    "Hmin_GeodLon_90",
    "Hmin_GeodLon_85",
    "Hmin_GeodLon_80",
    "Hmin_GeodLon_75",
    "Hmin_GeodLon_70",
    "Hmin_GeodLon_65",
    "Hmin_GeodLon_60",
    "Hmin_GeodLon_55",
    "Hmin_GeodLon_50",
    "Hmin_GeodLon_45",
    "Hmin_GeodLon_40",
    "Hmin_GeodLon_35",
    "Hmin_GeodLon_30",
    "Hmin_GeodLon_25",
    "Hmin_GeodLon_20",
    "Hmin_GeodLon_15",
    "Hmin_GeodLon_10",
    "Hmin_GeodLon_5",
    "Lstar_90",
    "Lstar_85",
    "Lstar_80",
    "Lstar_75",
    "Lstar_70",
    "Lstar_65",
    "Lstar_60",
    "Lstar_55",
    "Lstar_50",
    "Lstar_45",
    "Lstar_40",
    "Lstar_35",
    "Lstar_30",
    "Lstar_25",
    "Lstar_20",
    "Lstar_15",
    "Lstar_10",
    "Lstar_5",
    "DSType_90",
    "DSType_85",
    "DSType_80",
    "DSType_75",
    "DSType_70",
    "DSType_65",
    "DSType_60",
    "DSType_55",
    "DSType_50",
    "DSType_45",
    "DSType_40",
    "DSType_35",
    "DSType_30",
    "DSType_25",
    "DSType_20",
    "DSType_15",
    "DSType_10",
    "DSType_5",
    "L_90",
    "L_85",
    "L_80",
    "L_75",
    "L_70",
    "L_65",
    "L_60",
    "L_55",
    "L_50",
    "L_45",
    "L_40",
    "L_35",
    "L_30",
    "L_25",
    "L_20",
    "L_15",
    "L_10",
    "L_5",
    "Bm_90",
    "Bm_85",
    "Bm_80",
    "Bm_75",
    "Bm_70",
    "Bm_65",
    "Bm_60",
    "Bm_55",
    "Bm_50",
    "Bm_45",
    "Bm_40",
    "Bm_35",
    "Bm_30",
    "Bm_25",
    "Bm_20",
    "Bm_15",
    "Bm_10",
    "Bm_5",
    "I_90",
    "I_85",
    "I_80",
    "I_75",
    "I_70",
    "I_65",
    "I_60",
    "I_55",
    "I_50",
    "I_45",
    "I_40",
    "I_35",
    "I_30",
    "I_25",
    "I_20",
    "I_15",
    "I_10",
    "I_5",
    "K_90",
    "K_85",
    "K_80",
    "K_75",
    "K_70",
    "K_65",
    "K_60",
    "K_55",
    "K_50",
    "K_45",
    "K_40",
    "K_35",
    "K_30",
    "K_25",
    "K_20",
    "K_15",
    "K_10",
    "K_5",
    "Sb_90",
    "Sb_85",
    "Sb_80",
    "Sb_75",
    "Sb_70",
    "Sb_65",
    "Sb_60",
    "Sb_55",
    "Sb_50",
    "Sb_45",
    "Sb_40",
    "Sb_35",
    "Sb_30",
    "Sb_25",
    "Sb_20",
    "Sb_15",
    "Sb_10",
    "Sb_5",
    "Tb_90",
    "Tb_85",
    "Tb_80",
    "Tb_75",
    "Tb_70",
    "Tb_65",
    "Tb_60",
    "Tb_55",
    "Tb_50",
    "Tb_45",
    "Tb_40",
    "Tb_35",
    "Tb_30",
    "Tb_25",
    "Tb_20",
    "Tb_15",
    "Tb_10",
    "Tb_5",
    "Kappa_90",
    "Kappa_85",
    "Kappa_80",
    "Kappa_75",
    "Kappa_70",
    "Kappa_65",
    "Kappa_60",
    "Kappa_55",
    "Kappa_50",
    "Kappa_45",
    "Kappa_40",
    "Kappa_35",
    "Kappa_30",
    "Kappa_25",
    "Kappa_20",
    "Kappa_15",
    "Kappa_10",
    "Kappa_5",
]

name_278 = [
    "Time",
    "Date",
    "DOY",
    "UTC",
    "JulianDate",
    "GpsTime",
    "DipoleTiltAngle",
    "InOut",
    "Orbnumber",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "Rgeod_Lat",
    "Rgeod_Lon",
    "Rgeod_Height",
    "Rgsm_x",
    "Rgsm_y",
    "Rgsm_z",
    "Rsm_x",
    "Rsm_y",
    "Rsm_z",
    "Rgei_x",
    "Rgei_y",
    "Rgei_z",
    "Rgse_x",
    "Rgse_y",
    "Rgse_z",
    "CDMAG_MLAT",
    "CDMAG_MLON",
    "CDMAG_MLT",
    "CDMAG_R",
    "EDMAG_MLAT",
    "EDMAG_MLON",
    "EDMAG_MLT",
    "EDMAG_R",
    "IntModel",
    "ExtModel",
    "Kp",
    "Dst",
    "Bsc_gsm_x",
    "Bsc_gsm_y",
    "Bsc_gsm_z",
    "Bsc_gsm_mag",
    "FieldLineType",
    "S_sc_to_pfn",
    "S_sc_to_pfs",
    "S_pfs_to_Bmin",
    "S_Bmin_to_sc",
    "S_total",
    "d2B_ds2",
    "Sb0",
    "RadiusOfCurv",
    "Pfn_gsm_x",
    "Pfn_gsm_y",
    "Pfn_gsm_z",
    "Pfn_geo_x",
    "Pfn_geo_y",
    "Pfn_geo_z",
    "Pfn_geod_Lat",
    "Pfn_geod_Lon",
    "Pfn_geod_Height",
    "Pfn_CD_MLAT",
    "Pfn_CD_MLON",
    "Pfn_CD_MLT",
    "Pfn_ED_MLAT",
    "Pfn_ED_MLON",
    "Pfn_ED_MLT",
    "Bfn_geo_x",
    "Bfn_geo_y",
    "Bfn_geo_z",
    "Bfn_geo_mag",
    "Bfn_gsm_x",
    "Bfn_gsm_y",
    "Bfn_gsm_z",
    "Bfn_gsm_mag",
    "Loss_Cone_Alpha_n",
    "Pfs_gsm_x",
    "Pfs_gsm_y",
    "Pfs_gsm_z",
    "Pfs_geo_x",
    "Pfs_geo_y",
    "Pfs_geo_z",
    "Pfs_geod_Lat",
    "Pfs_geod_Lon",
    "Pfs_geod_Height",
    "Pfs_CD_MLAT",
    "Pfs_CD_MLON",
    "Pfs_CD_MLT",
    "Pfs_ED_MLAT",
    "Pfs_ED_MLON",
    "Pfs_ED_MLT",
    "Bfs_geo_x",
    "Bfs_geo_y",
    "Bfs_geo_z",
    "Bfs_geo_mag",
    "Bfs_gsm_x",
    "Bfs_gsm_y",
    "Bfs_gsm_z",
    "Bfs_gsm_mag",
    "Loss_Cone_Alpha_s",
    "Pmin_gsm_x",
    "Pmin_gsm_y",
    "Pmin_gsm_z",
    "Bmin_gsm_x",
    "Bmin_gsm_y",
    "Bmin_gsm_z",
    "Bmin_gsm_mag",
    "Lsimple",
    "InvLat",
    "Lm_eq",
    "InvLat_eq",
    "BoverBeq",
    "MlatFromBoverBeq",
    "M_used",
    "M_ref",
    "M_igrf",
    "Lstar_90",
    "Lstar_85",
    "Lstar_80",
    "Lstar_75",
    "Lstar_70",
    "Lstar_65",
    "Lstar_60",
    "Lstar_55",
    "Lstar_50",
    "Lstar_45",
    "Lstar_40",
    "Lstar_35",
    "Lstar_30",
    "Lstar_25",
    "Lstar_20",
    "Lstar_15",
    "Lstar_10",
    "Lstar_5",
    "DSType_90",
    "DSType_85",
    "DSType_80",
    "DSType_75",
    "DSType_70",
    "DSType_65",
    "DSType_60",
    "DSType_55",
    "DSType_50",
    "DSType_45",
    "DSType_40",
    "DSType_35",
    "DSType_30",
    "DSType_25",
    "DSType_20",
    "DSType_15",
    "DSType_10",
    "DSType_5",
    "L_90",
    "L_85",
    "L_80",
    "L_75",
    "L_70",
    "L_65",
    "L_60",
    "L_55",
    "L_50",
    "L_45",
    "L_40",
    "L_35",
    "L_30",
    "L_25",
    "L_20",
    "L_15",
    "L_10",
    "L_5",
    "Bm_90",
    "Bm_85",
    "Bm_80",
    "Bm_75",
    "Bm_70",
    "Bm_65",
    "Bm_60",
    "Bm_55",
    "Bm_50",
    "Bm_45",
    "Bm_40",
    "Bm_35",
    "Bm_30",
    "Bm_25",
    "Bm_20",
    "Bm_15",
    "Bm_10",
    "Bm_5",
    "I_90",
    "I_85",
    "I_80",
    "I_75",
    "I_70",
    "I_65",
    "I_60",
    "I_55",
    "I_50",
    "I_45",
    "I_40",
    "I_35",
    "I_30",
    "I_25",
    "I_20",
    "I_15",
    "I_10",
    "I_5",
    "K_90",
    "K_85",
    "K_80",
    "K_75",
    "K_70",
    "K_65",
    "K_60",
    "K_55",
    "K_50",
    "K_45",
    "K_40",
    "K_35",
    "K_30",
    "K_25",
    "K_20",
    "K_15",
    "K_10",
    "K_5",
    "Sb_90",
    "Sb_85",
    "Sb_80",
    "Sb_75",
    "Sb_70",
    "Sb_65",
    "Sb_60",
    "Sb_55",
    "Sb_50",
    "Sb_45",
    "Sb_40",
    "Sb_35",
    "Sb_30",
    "Sb_25",
    "Sb_20",
    "Sb_15",
    "Sb_10",
    "Sb_5",
    "Tb_90",
    "Tb_85",
    "Tb_80",
    "Tb_75",
    "Tb_70",
    "Tb_65",
    "Tb_60",
    "Tb_55",
    "Tb_50",
    "Tb_45",
    "Tb_40",
    "Tb_35",
    "Tb_30",
    "Tb_25",
    "Tb_20",
    "Tb_15",
    "Tb_10",
    "Tb_5",
    "Kappa_90",
    "Kappa_85",
    "Kappa_80",
    "Kappa_75",
    "Kappa_70",
    "Kappa_65",
    "Kappa_60",
    "Kappa_55",
    "Kappa_50",
    "Kappa_45",
    "Kappa_40",
    "Kappa_35",
    "Kappa_30",
    "Kappa_25",
    "Kappa_20",
    "Kappa_15",
    "Kappa_10",
    "Kappa_5",
]

h5_mag_name = [
    "Alpha",
    "ApogeePosGeod",
    "ApogeeTimes",
    "AscendPosGeod",
    "AscendTimes",
    "Bfn_geo",
    "Bfn_gsm",
    "Bfs_geo",
    "Bfs_gsm",
    "Bm",
    "Bmin_gsm",
    "BoverBeq",
    "Bsc_gsm",
    "CDMAG_MLAT",
    "CDMAG_MLON",
    "CDMAG_MLT",
    "CDMAG_R",
    "Date",
    "DipoleTiltAngle",
    "Doy",
    "DriftShellType",
    "Dst",
    "EDMAG_MLAT",
    "EDMAG_MLON",
    "EDMAG_MLT",
    "EDMAG_R",
    "ExtModel",
    "FieldLineType",
    "GpsTime",
    "I",
    "InOut",
    "IntModel",
    "InvLat",
    "InvLat_eq",
    "IsoTime",
    "JulianDate",
    "K",
    "Kappa",
    "Kp",
    "L",
    "Lm_eq",
    "Loss_Cone_Alpha_n",
    "Loss_Cone_Alpha_s",
    "Lsimple",
    "Lstar",
    "M_igrf",
    "M_ref",
    "M_used",
    "MlatFromBoverBeq",
    "OrbitNumber",
    "PerigeePosGeod",
    "PerigeeTimes",
    "Pfn_CD_MLAT",
    "Pfn_CD_MLON",
    "Pfn_CD_MLT",
    "Pfn_ED_MLAT",
    "Pfn_ED_MLON",
    "Pfn_ED_MLT",
    "Pfn_geo",
    "Pfn_geod_Height",
    "Pfn_geod_LatLon",
    "Pfn_gsm",
    "Pfs_CD_MLAT",
    "Pfs_CD_MLON",
    "Pfs_CD_MLT",
    "Pfs_ED_MLAT",
    "Pfs_ED_MLON",
    "Pfs_ED_MLT",
    "Pfs_geo",
    "Pfs_geod_Height",
    "Pfs_geod_LatLon",
    "Pfs_gsm",
    "Pmin_gsm",
    "RadiusOfCurv",
    "Rgei",
    "Rgeo",
    "Rgeod_Height",
    "Rgeod_LatLon",
    "Rgse",
    "Rgsm",
    "Rsm",
    "S_Bmin_to_sc",
    "S_pfs_to_Bmin",
    "S_sc_to_pfn",
    "S_sc_to_pfs",
    "S_total",
    "Sb",
    "Sb0",
    "Tb",
    "UTC",
    "d2B_ds2",
]


def getstartline(file_name):
    """
    Find the line number of the header.
    """

    with open(file_name, "rb") as f:
        line_total_num = sum(1 for line in f)

    start_line = 0
    for i in range(920, line_total_num):
        line_data = linecache.getline(file_name, i)

        # get the first 10
        if line_data[0:10] == "#  +------":
            start_line = i
            break
    if start_line == 0:
        raise ValueError("the txt file is wrong")
    return start_line


def magephem_read_txt(out_files, varnames=None, notplot=False, prefix="", suffix=""):
    """
    Read the txt file and return the data as a dictionary or save to tplot variables.
    """

    if (varnames is None) or (len(varnames) == 0):
        varnames2load = ["Time", "Lsimple", "Lm_eq", "CDMAG_MLAT", "CDMAG_MLON", "CDMAG_MLT"]
    else:
        varnames2load = varnames

    if "Time" not in varnames2load:
        varnames2load.append("Time")

    data_all = []

    for ifile, file in enumerate(out_files):
        start_line = getstartline(file)
        # Read the data using numpy
        file_pd = pd.read_csv(file, skiprows=(start_line + 2), sep=r"\s+", header=None)

        # If the time in the last row is 00:00:00, which duplicates the start of the next day, remove the last row
        if (ifile != len(out_files) - 1) and ("T00:00:00.0000Z" in file_pd.iloc[-1, 0]):
            file_pd = file_pd.iloc[:-1, :]

        if len(file_pd.shape) == 2:
            if file_pd.shape[1] == 278:
                file_pd.columns = name_278
            elif file_pd.shape[1] == 332:
                file_pd.columns = name_332
            else:
                raise ValueError(f"This date file has some problem with the number of columns {file_pd.shape}!! \n {file}")
        else:
            raise ValueError(f"This date file has some problem with only one column!! \n {file}")

        # choose the columns that
        for name in varnames2load:
            if name not in file_pd.columns:
                raise ValueError(f"The variable {name} is not in the txt file. \n {file}")

        select_pd = file_pd[varnames2load]
        select_pd.loc[:, "Time"] = pd.to_datetime(select_pd["Time"])
        data_all.append(select_pd)

    data_all = pd.concat(data_all, ignore_index=True)
    final_dict = data_all.to_dict(orient="list")
    for key in final_dict.keys():
        final_dict[key] = np.array(final_dict[key])

    if notplot:
        return final_dict
    else:
        for key in final_dict.keys():
            if key == "Time":
                continue
            store_data(prefix + key + suffix, data={"x": final_dict["Time"], "y": final_dict[key]})
        return list(final_dict.keys())


def magephem_read_h5(out_files, varnames=None, notplot=False, prefix="", suffix=""):
    """
    Read the h5 file and return the data as a dictionary or save to tplot variables.
    """

    if (varnames is None) or (len(varnames) == 0):
        varnames = ["IsoTime", "Lsimple", "Lm_eq", "CDMAG_MLAT", "CDMAG_MLON", "CDMAG_MLT"]

    data_all = []

    for ifile, file in enumerate(out_files):
        # open the h5 file
        f = h5py.File(file, "r")
        # ignore the foot of data
        date_array = np.char.decode(np.array(f.get("IsoTime")))
        time_dt64 = pd.to_datetime(date_array)
        file_pd = pd.DataFrame(columns=varnames)
        file_pd["IsoTime"] = time_dt64
        for name in varnames:
            if name == "IsoTime":
                continue
            if name not in h5_mag_name:
                raise ValueError(f"The variable {name} is not in the h5 file. \n {h5_mag_name}")
            else:
                data_name = np.array(f.get(name))
                file_pd[name] = data_name

        # If the time in the last row is 00:00:00, which duplicates the start of the next day, remove the last row
        if (ifile != len(out_files) - 1) and ("00:00:00" in date_array[-1]):
            file_pd = file_pd.iloc[:-1, :]

        data_all.append(file_pd)

    data_all = pd.concat(data_all, ignore_index=True)
    data_all = data_all.rename(columns={"IsoTime": "Time"})
    final_dict = data_all.to_dict(orient="list")

    for key in final_dict.keys():
        final_dict[key] = np.array(final_dict[key])

    if notplot:
        return final_dict
    else:
        for key in final_dict.keys():
            if key == "Time":
                continue
            store_data(prefix + key + suffix, data={"x": final_dict["Time"], "y": final_dict[key]})
        return list(final_dict.keys())
