"""
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",
    "Julian Date",
    "GPS Time",
    "Titleangle",
    "InOut",
    "Orbnumber",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "GeodLat",
    "GeodLon",
    "GeodHeight",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Xsm",
    "Ysm",
    "Zsm",
    "Xgei",
    "Ygei",
    "Zgei",
    "Xgse",
    "Ygse",
    "Zgse",
    "CD_MLAT",
    "CD_MLON",
    "CD_MLT",
    "CD_R",
    "ED_MLAT",
    "ED_MLON",
    "ED_MLT",
    "ED_R",
    "Int Model",
    "Ext Model",
    "Kp",
    "Dst",
    "Bsc_x",
    "Bsc_y",
    "Bsc_z",
    "Bsc",
    "Field Line Type",
    "S_sc_to_pfn",
    "S_sc_to_pfs",
    "S_pfs_to_Bmin",
    "S_sc_to_Bmin",
    "S_total",
    "d2B_ds2",
    "Sb0",
    "RadOfCurv",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "GeodLat",
    "GeodLon",
    "GeodHeight",
    "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",
    "Bfn_x",
    "Bfn_y",
    "Bfn_z",
    "Bfn",
    "Alpha_LC_N",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "GeodLat",
    "GeodLon",
    "GeodHeight",
    "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",
    "Bfs_x",
    "Bfs_y",
    "Bfs_z",
    "Bfs",
    "Alpha_LC_S",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Bmin_x",
    "Bmin_y",
    "Bmin_z",
    "Bmin",
    "Lsimple",
    "InvLat",
    "Lm_eq",
    "InvLatEq",
    "BoverBeq",
    "MlatBBeq",
    "M_used",
    "M_ref",
    "M_igrf",
    "Hmin0",
    "Hmin1",
    "Hmin2",
    "Hmin3",
    "Hmin4",
    "Hmin5",
    "Hmin6",
    "Hmin7",
    "Hmin8",
    "Hmin9",
    "Hmin10",
    "Hmin11",
    "Hmin12",
    "Hmin13",
    "Hmin14",
    "Hmin15",
    "Hmin16",
    "Hmin17",
    "Hmin_GeodLat0",
    "Hmin_GeodLat1",
    "Hmin_GeodLat2",
    "Hmin_GeodLat3",
    "Hmin_GeodLat4",
    "Hmin_GeodLat5",
    "Hmin_GeodLat6",
    "Hmin_GeodLat7",
    "Hmin_GeodLat8",
    "Hmin_GeodLat9",
    "Hmin_GeodLat10",
    "Hmin_GeodLat11",
    "Hmin_GeodLat12",
    "Hmin_GeodLat13",
    "Hmin_GeodLat14",
    "Hmin_GeodLat15",
    "Hmin_GeodLat16",
    "Hmin_GeodLat17",
    "Hmin_GeodLon0",
    "Hmin_GeodLon1",
    "Hmin_GeodLon2",
    "Hmin_GeodLon3",
    "Hmin_GeodLon4",
    "Hmin_GeodLon5",
    "Hmin_GeodLon6",
    "Hmin_GeodLon7",
    "Hmin_GeodLon8",
    "Hmin_GeodLon9",
    "Hmin_GeodLon10",
    "Hmin_GeodLon11",
    "Hmin_GeodLon12",
    "Hmin_GeodLon13",
    "Hmin_GeodLon14",
    "Hmin_GeodLon15",
    "Hmin_GeodLon16",
    "Hmin_GeodLon17",
    "L*0",
    "L*1",
    "L*2",
    "L*3",
    "L*4",
    "L*5",
    "L*6",
    "L*7",
    "L*8",
    "L*9",
    "L*10",
    "L*11",
    "L*12",
    "L*13",
    "L*14",
    "L*15",
    "L*16",
    "L*17",
    "DSType0",
    "DSType1",
    "DSType2",
    "DSType3",
    "DSType4",
    "DSType5",
    "DSType6",
    "DSType7",
    "DSType8",
    "DSType9",
    "DSType10",
    "DSType11",
    "DSType12",
    "DSType13",
    "DSType14",
    "DSType15",
    "DSType16",
    "DSType17",
    "L0",
    "L1",
    "L2",
    "L3",
    "L4",
    "L5",
    "L6",
    "L7",
    "L8",
    "L9",
    "L10",
    "L11",
    "L12",
    "L13",
    "L14",
    "L15",
    "L16",
    "L17",
    "Bm0",
    "Bm1",
    "Bm2",
    "Bm3",
    "Bm4",
    "Bm5",
    "Bm6",
    "Bm7",
    "Bm8",
    "Bm9",
    "Bm10",
    "Bm11",
    "Bm12",
    "Bm13",
    "Bm14",
    "Bm15",
    "Bm16",
    "Bm17",
    "I0",
    "I1",
    "I2",
    "I3",
    "I4",
    "I5",
    "I6",
    "I7",
    "I8",
    "I9",
    "I10",
    "I11",
    "I12",
    "I13",
    "I14",
    "I15",
    "I16",
    "I17",
    "K0",
    "K1",
    "K2",
    "K3",
    "K4",
    "K5",
    "K6",
    "K7",
    "K8",
    "K9",
    "K10",
    "K11",
    "K12",
    "K13",
    "K14",
    "K15",
    "K16",
    "K17",
    "Sb0",
    "Sb1",
    "Sb2",
    "Sb3",
    "Sb4",
    "Sb5",
    "Sb6",
    "Sb7",
    "Sb8",
    "Sb9",
    "Sb10",
    "Sb11",
    "Sb12",
    "Sb13",
    "Sb14",
    "Sb15",
    "Sb16",
    "Sb17",
    "Tb0",
    "Tb1",
    "Tb2",
    "Tb3",
    "Tb4",
    "Tb5",
    "Tb6",
    "Tb7",
    "Tb8",
    "Tb9",
    "Tb10",
    "Tb11",
    "Tb12",
    "Tb13",
    "Tb14",
    "Tb15",
    "Tb16",
    "Tb17",
    "Kappa0",
    "Kappa1",
    "Kappa2",
    "Kappa3",
    "Kappa4",
    "Kappa5",
    "Kappa6",
    "Kappa7",
    "Kappa8",
    "Kappa9",
    "Kappa10",
    "Kappa11",
    "Kappa12",
    "Kappa13",
    "Kappa14",
    "Kappa15",
    "Kappa16",
    "Kappa17",
]

name_278 = [
    "Time",
    "Date",
    "DOY",
    "UTC",
    "Julian Date",
    "GPS Time",
    "Titleangle",
    "InOut",
    "Orbnumber",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "GeodLat",
    "GeodLon",
    "GeodHeight",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Xsm",
    "Ysm",
    "Zsm",
    "Xgei",
    "Ygei",
    "Zgei",
    "Xgse",
    "Ygse",
    "Zgse",
    "CD_MLAT",
    "CD_MLON",
    "CD_MLT",
    "CD_R",
    "ED_MLAT",
    "ED_MLON",
    "ED_MLT",
    "ED_R",
    "Int Model",
    "Ext Model",
    "Kp",
    "Dst",
    "Bsc_x",
    "Bsc_y",
    "Bsc_z",
    "Bsc",
    "Field Line Type",
    "S_sc_to_pfn",
    "S_sc_to_pfs",
    "S_pfs_to_Bmin",
    "S_sc_to_Bmin",
    "S_total",
    "d2B_ds2",
    "Sb0",
    "RadOfCurv",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "GeodLat",
    "GeodLon",
    "GeodHeight",
    "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",
    "Bfn_x",
    "Bfn_y",
    "Bfn_z",
    "Bfn",
    "Alpha_LC_N",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Xgeo",
    "Ygeo",
    "Zgeo",
    "GeodLat",
    "GeodLon",
    "GeodHeight",
    "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",
    "Bfs_x",
    "Bfs_y",
    "Bfs_z",
    "Bfs",
    "Alpha_LC_S",
    "Xgsm",
    "Ygsm",
    "Zgsm",
    "Bmin_x",
    "Bmin_y",
    "Bmin_z",
    "Bmin",
    "Lsimple",
    "InvLat",
    "Lm_eq",
    "InvLatEq",
    "BoverBeq",
    "MlatBBeq",
    "M_used",
    "M_ref",
    "M_igrf",
    "L*0",
    "L*1",
    "L*2",
    "L*3",
    "L*4",
    "L*5",
    "L*6",
    "L*7",
    "L*8",
    "L*9",
    "L*10",
    "L*11",
    "L*12",
    "L*13",
    "L*14",
    "L*15",
    "L*16",
    "L*17",
    "DSType0",
    "DSType1",
    "DSType2",
    "DSType3",
    "DSType4",
    "DSType5",
    "DSType6",
    "DSType7",
    "DSType8",
    "DSType9",
    "DSType10",
    "DSType11",
    "DSType12",
    "DSType13",
    "DSType14",
    "DSType15",
    "DSType16",
    "DSType17",
    "L0",
    "L1",
    "L2",
    "L3",
    "L4",
    "L5",
    "L6",
    "L7",
    "L8",
    "L9",
    "L10",
    "L11",
    "L12",
    "L13",
    "L14",
    "L15",
    "L16",
    "L17",
    "Bm0",
    "Bm1",
    "Bm2",
    "Bm3",
    "Bm4",
    "Bm5",
    "Bm6",
    "Bm7",
    "Bm8",
    "Bm9",
    "Bm10",
    "Bm11",
    "Bm12",
    "Bm13",
    "Bm14",
    "Bm15",
    "Bm16",
    "Bm17",
    "I0",
    "I1",
    "I2",
    "I3",
    "I4",
    "I5",
    "I6",
    "I7",
    "I8",
    "I9",
    "I10",
    "I11",
    "I12",
    "I13",
    "I14",
    "I15",
    "I16",
    "I17",
    "K0",
    "K1",
    "K2",
    "K3",
    "K4",
    "K5",
    "K6",
    "K7",
    "K8",
    "K9",
    "K10",
    "K11",
    "K12",
    "K13",
    "K14",
    "K15",
    "K16",
    "K17",
    "Sb0",
    "Sb1",
    "Sb2",
    "Sb3",
    "Sb4",
    "Sb5",
    "Sb6",
    "Sb7",
    "Sb8",
    "Sb9",
    "Sb10",
    "Sb11",
    "Sb12",
    "Sb13",
    "Sb14",
    "Sb15",
    "Sb16",
    "Sb17",
    "Tb0",
    "Tb1",
    "Tb2",
    "Tb3",
    "Tb4",
    "Tb5",
    "Tb6",
    "Tb7",
    "Tb8",
    "Tb9",
    "Tb10",
    "Tb11",
    "Tb12",
    "Tb13",
    "Tb14",
    "Tb15",
    "Tb16",
    "Tb17",
    "Kappa0",
    "Kappa1",
    "Kappa2",
    "Kappa3",
    "Kappa4",
    "Kappa5",
    "Kappa6",
    "Kappa7",
    "Kappa8",
    "Kappa9",
    "Kappa10",
    "Kappa11",
    "Kappa12",
    "Kappa13",
    "Kappa14",
    "Kappa15",
    "Kappa16",
    "Kappa17",
]

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", "CD_MLAT", "CD_MLON", "CD_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
        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 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 final_dict.keys()
