import os
import glob

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# ==========================
# User parameters
# ==========================
# Read all files matching this pattern in the same directory as this script
file_pattern = "dBzdt_*.txt"

label_fontsize = 18
tick_fontsize = 15
legend_fontsize = 15
linewidth = 2.5

xmin = 1e-6
xmax = 1e-1
ymin = None
ymax = None

use_abs = True

figsize = (8, 6)

savefig = True
save_name = "TEM_decay_curve.png"
dpi = 600

# ==========================
# Colors
# ==========================
# Fixed-order categorical palette (colorblind-safe, validated).
# For more than 8 series, a sequential blue ramp is used instead of cycling.
PALETTE = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100",
           "#e87ba4", "#008300", "#4a3aa7", "#e34948"]
RAMP_HEX = ["#86b6ef", "#5598e7", "#3987e5", "#2a78d6",
            "#256abf", "#1c5cab", "#184f95", "#104281"]


def series_color(i, n):
    """Color for the i-th series of n (fixed order, never cycled)."""
    if n <= len(PALETTE):
        return PALETTE[i]
    cmap = LinearSegmentedColormap.from_list("seq_blue", RAMP_HEX)
    return cmap(i / (n - 1))


# ==========================
# Read data
# ==========================
script_dir = os.path.dirname(os.path.abspath(__file__))
files = sorted(glob.glob(os.path.join(script_dir, file_pattern)))

if not files:
    raise FileNotFoundError(
        f"No files matching '{file_pattern}' in {script_dir}")

if use_abs:
    ylabel = r"$|dB_z/dt|$ (V/A)"
else:
    ylabel = r"$dB_z/dt$ (V/A)"

plt.figure(figsize=figsize)

for i, filename in enumerate(files):
    # Header: line 1 = point name, line 2 = receiver coordinates (x, y, z)
    with open(filename) as f:
        point_name = f.readline().strip()
        coords = [float(v) for v in f.readline().split()]
    x, y, z = coords
    label = f"{point_name}  ({x:g}, {y:g}, {z:g}) m"

    data = np.loadtxt(filename, skiprows=2)

    time = data[:, 1]
    dbdt = data[:, 2]

    if use_abs:
        dbdt_plot = np.abs(dbdt)
    else:
        dbdt_plot = dbdt

    mask = (time > 0) & (dbdt_plot > 0)

    plt.loglog(time[mask], dbdt_plot[mask], linewidth=linewidth,
               color=series_color(i, len(files)), label=label)

# ==========================
# Plot
# ==========================
plt.xlabel("Time (s)", fontsize=label_fontsize)
plt.ylabel(ylabel, fontsize=label_fontsize)

plt.xticks(fontsize=tick_fontsize)
plt.yticks(fontsize=tick_fontsize)

plt.grid(True, which="both", linestyle="--", alpha=0.4)

if xmin is not None or xmax is not None:
    plt.xlim(xmin, xmax)

if ymin is not None or ymax is not None:
    plt.ylim(ymin, ymax)

plt.legend(fontsize=legend_fontsize)
plt.tight_layout()

if savefig:
    plt.savefig(os.path.join(script_dir, save_name), dpi=dpi)

plt.show()