# -*- coding: utf-8 -*-
"""
Created on Mon Oct 19 13:41:09 2026

@author: chuffard
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize



def sliding_ccf(ts_lead, ts_lag, win_width=60,
                lag_max=30, plot_results=True, dates=None):

    """
    Sliding-window cross-correlation.

    Parameters
    ----------
    ts_lead : array-like
        Leading variable (e.g., POCmodel_mg).
    ts_lag : array-like
        Lagging variable (e.g., total_area_flux).
    win_width : int
        Number of records in each window.
    lag_max : int
        Maximum absolute lag to examine.
    plot_results : bool
        Whether to plot the lag and correlation profiles.
    dates : array-like, optional
        Dates for the x-axis.

    Returns
    -------
    DataFrame with columns:
        MaxNegLag
        MaxNegCorrelation
    """

    # Convert inputs to numeric arrays
    lead = np.asarray(ts_lead, dtype=float)
    lag = np.asarray(ts_lag, dtype=float)

    # Check inputs
    if len(lead) != len(lag):
        raise ValueError(
            "Lead and lag time series must have equal lengths."
        )

    n = len(lead)

    if win_width < 2:
        raise ValueError("win_width must be at least 2.")

    if n < 2:
        raise ValueError("At least two observations are required.")

    if win_width > n:
        print("Window exceeds data length; using full data length.")
        win_width = n

    if lag_max >= win_width:
        print("lag_max exceeds window width; reducing it.")
        lag_max = win_width - 1

    if lag_max < 0:
        raise ValueError("lag_max cannot be negative.")

    half = win_width // 2

    # Match the original R window alignment.
    # For even windows, the center falls between two records.
    start_pos = half - 1 if win_width % 2 == 0 else half

    # Initialize outputs, retaining one result per original record.
    max_corr = np.full(n, np.nan)
    max_lag = np.full(n, np.nan)

    # Calculate correlation at a particular lag.
    # This convention matches R's ccf(x, y) lag definition:
    # lag k compares lead[t+k] with lag[t].
    def lag_correlation(x, y, k):

        if k < 0:
            x_pair = x[:k]
            y_pair = y[-k:]
        elif k > 0:
            x_pair = x[k:]
            y_pair = y[:-k]
        else:
            x_pair = x
            y_pair = y

        # Pairwise removal of missing/non-finite observations
        valid = np.isfinite(x_pair) & np.isfinite(y_pair)
        x_pair = x_pair[valid]
        y_pair = y_pair[valid]

        if len(x_pair) < 2:
            return np.nan

        # A correlation is undefined for a constant series.
        if np.std(x_pair) == 0 or np.std(y_pair) == 0:
            return np.nan

        return np.corrcoef(x_pair, y_pair)[0, 1]

    # Only evaluate negative lags and zero, as in the R script.
    lags_to_test = np.arange(-lag_max, 1)

    # Sliding-window loop
    for i in range(start_pos, n - half):

        if win_width % 2 == 0:
            window_start = i - half + 1
        else:
            window_start = i - half

        window_end = window_start + win_width

        x = lead[window_start:window_end]
        y = lag[window_start:window_end]

        correlations = np.array([
            lag_correlation(x, y, k)
            for k in lags_to_test
        ])

        valid = np.isfinite(correlations)

        if not valid.any():
            continue

        # Select the largest ABSOLUTE correlation,
        # but retain its original sign in the output.
        valid_indices = np.flatnonzero(valid)
        best_idx = valid_indices[
            np.argmax(np.abs(correlations[valid]))
        ]

        max_corr[i] = correlations[best_idx]
        max_lag[i] = lags_to_test[best_idx]

    # Return one row per original observation.
    results = pd.DataFrame({
        "MaxNegLag": max_lag,
        "MaxNegCorrelation": max_corr
    })

    if plot_results:

        x_axis = (
            np.asarray(dates)
            if dates is not None
            else np.arange(n)
        )

        valid_results = (
            np.isfinite(max_lag) & np.isfinite(max_corr)
        )

        # Only color points where the selected correlation is positive
        positive_corr = valid_results & (max_corr > 0)

        # Plot 1: profile colored by lag (unchanged)
        fig, ax = plt.subplots(figsize=(13, 5))

        ax.plot(
            x_axis, lag,
            color="lightgray",
            linewidth=1.5,
            zorder=1
        )

        sc = ax.scatter(
            np.asarray(x_axis)[positive_corr],
            lag[positive_corr],
            c=max_lag[positive_corr],
            cmap="hot_r",
            norm=Normalize(vmin=-lag_max, vmax=0),
            s=18,
            zorder=2
        )
        
        
        # sc = ax.scatter(
        #     np.asarray(x_axis)[positive_corr],
        #     lag[positive_corr],
        #     c=max_corr[positive_corr],
        #     cmap="coolwarm",
        #     norm=Normalize(vmin=0, vmax=1),
        #     s=18,
        #     zorder=2
        # )

        ax.set_title(
            f"Lagging Variable Profile Colored by Lag "
            f"(window = {win_width} records)"
        )
        ax.set_ylabel("Lagging variable")
        ax.set_xlabel("Date" if dates is not None else "Record")
        fig.colorbar(sc, ax=ax, label="Selected lag (records)")
        fig.tight_layout()
        plt.show()

        # Plot 2: profile colored by POSITIVE correlation only
        fig, ax = plt.subplots(figsize=(13, 5))

        ax.plot(
            x_axis, lag,
            color="lightgray",
            linewidth=1.5,
            zorder=1
        )

        sc = ax.scatter(
            np.asarray(x_axis)[positive_corr],
            lag[positive_corr],
            c=max_corr[positive_corr],
            cmap="coolwarm",
            norm=Normalize(vmin=0, vmax=1),
            s=18,
            zorder=2
        )

        ax.set_title(
            f"Lagging Variable Profile Colored by Positive Correlation "
            f"(window = {win_width} records)"
        )
        ax.set_ylabel("Lagging variable")
        ax.set_xlabel("Date" if dates is not None else "Record")

        if positive_corr.any():
            fig.colorbar(
                sc, ax=ax,
                label="Positive correlation"
            )

        fig.tight_layout()
        plt.show()

        
        # Plot 3: Leading and lagging variables with two y-axes
        fig, ax1 = plt.subplots(figsize=(13, 5))

       # Left y-axis: leading variable
        line1 = ax1.plot(
           x_axis,
           lead,
           color="tab:blue",
           linewidth=1.5,
           label="POCmodel_mg"
       )
        ax1.set_xlabel("Date" if dates is not None else "Record")
        ax1.set_ylabel("POCmodel_mg", color="tab:blue")
        ax1.tick_params(axis="y", labelcolor="tab:blue")

       # Right y-axis: lagging variable
        ax2 = ax1.twinx()

        line2 = ax2.plot(
           x_axis,
           lag,
           color="tab:red",
           linewidth=1.5,
           label="total_area_flux"
       )
        ax2.set_ylabel("total_area_flux", color="tab:red")
        ax2.tick_params(axis="y", labelcolor="tab:red")

       # Combined legend
        lines = line1 + line2
        labels = [line.get_label() for line in lines]
        ax1.legend(lines, labels, loc="best")

        ax1.set_title(
           "Leading and Lagging Variables Through Time"
       )

        fig.tight_layout()
        plt.show()

    return results


#now run it
slidingdata = SES_OffShore_SINKER[['date', 'total_area_flux', 'POCmodel_mg_daily_avg']]

#slidingdatadaily = slidingdata.groupby("date").transform("mean")
slidingdatadaily = slidingdata.resample('D').mean()

myData = slidingdatadaily

# Use your actual date column name if one exists.
# Replace "date" below with its name.
myData["date"] = pd.to_datetime(myData["date"])

mySlidingCCF = sliding_ccf(
    ts_lead=myData["POCmodel_mg_daily_avg"],
    ts_lag=myData["total_area_flux"],
    win_width=60,
    lag_max=30,
    plot_results=True,
    dates=myData["date"]
)