import pandas as pd
import numpy as np
import io
import pyTMD
from pyTMD.io.constituents import constituents
from eo_tides.utils import _standardise_time
from eo_tides.model import model_tides


def extract_constants_austides(file_path):
    """
    Parse harmonic constants from an AusTides harmonic constants file.
    """
    # Open file containing harmonic constants
    with open(file_path, "r") as f:
        text = f.read()

    # Break into sections
    sections = text.split("\n\n")

    # Find section containing LAT and LONG
    coord_section = next(s for s in sections if "LAT" in s and "LONG" in s)
    parts = coord_section.split()

    # Convert to decimal degrees (assuming parts will be:
    # ['Time', 'Zone', '−0800', 'LAT', '-18', '0', 'LONG', '122', '13'])
    lat = float(parts[4]) + float(parts[5]) / 60
    lon = float(parts[7]) + float(parts[8]) / 60

    # Find section containing constants
    harmonic_idx = sections.index("HARMONIC CONSTANTS")
    data_text = sections[harmonic_idx + 2]

    # Read constants
    constant_df = pd.read_csv(
        io.StringIO(data_text),
        sep="\s+",
        names=["c", "amp", "ph"],
    )

    # Rename constants to match pyTMD conventions
    c = [constituents.parse(i) for i in constant_df.c]

    # Reshape extracted constants and convert to
    # masked arrays to for pyTMD compatability
    amp = np.expand_dims(constant_df.amp, axis=0)
    ph = np.expand_dims(constant_df.ph, axis=0)
    amp = np.ma.array(amp, mask=False)
    ph = np.ma.array(ph, mask=False)

    # Return constants and lat/lon coords
    return amp, ph, c, lon, lat


# Input params
time = pd.date_range("2020-01", "2020-02", freq="3h")
austides_path = "/home/jovyan/Robbi/eo-tides/tests/data/63564_HCons.txt"

# Extract AusTides constants
amp, ph, c, lon, lat = extract_constants_austides(austides_path)

# Model tides with `eo-tides` and a global model
tides_df = model_tides(
    x=lon,
    y=lat,
    time=time,
    model="EOT20",
    directory="/var/share/tide_models/",
    output_format="wide",
)

# Standardise time to numpy datetime64 format
time = _standardise_time(time)

# Convert datetime
timescale = pyTMD.time.timescale().from_datetime(time.flatten())

# Calculate complex phase in radians for Euler's
cph = -1j * ph * np.pi / 180.0

# Calculate constituent oscillation
hc = amp * np.exp(cph)

# Convert datetime
deltat = np.zeros_like(timescale.tt_ut1)
t = timescale.tide

# Create arrays to hold outputs
tide = np.ma.zeros((len(t)), fill_value=np.nan)
tide.mask = np.any(hc.mask, axis=1)

# Predict tides
tides_austide = pyTMD.predict.time_series(
    t,
    hc,
    c,
    deltat=deltat.flatten(),
)

# Add modelled AusTides to original dataframe
tides_df["AusTides"] = tides_austide

# Plot
tides_df.droplevel(["x", "y"]).plot()
tides_df.droplevel(["x", "y"]).plot.scatter("AusTides", "EOT20")