import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime

plt.style.use("dark_background")

# USER PARAMETERS
TICKER = "ASML.AS"           # Stock symbol to fetch data from Yahoo Finance
START_DATE = "2023-01-01"    # Start date of historical data
END_DATE   = "2025-06-01"    # End date of historical data

# --- Main Logic Mode ---
METHOD = "Rolling"           # "Static" => regression per fixed window. "Rolling" => regression updates bar-by-bar

# --- Anchor Logic ---
ANCHOR_TYPE = "Periodic"     # "First Bar" => anchor from start. "Periodic" => reset based on time
ANCHOR_RES  = "weekly"       # Defines how frequently the anchor resets ("daily", "weekly", "monthly", etc.)

# --- Price Column ---
SOURCE_COL = "Close"         # Which column to use for regression (e.g., "Close", "Adj Close")

# --- Band Settings ---
MULT = 1.0                   # Multiplier for standard deviation bands (higher = wider bands)
SHOW_BANDS = True            # If True, draw upper/lower deviation bands around regression line

# --- Color Config ---
BICOLOR    = True            # If True, use separate colors based on slope sign (up/down)
MONO_COLOR = "#FF5D00"       # If not bicolor, this color will be used for all lines
COL_A = "#089981"            # Slope >= 0 color
COL_B = "#F23645"            # Slope < 0 color

# --- Visual Enhancements ---
GRADIENT = True              # If True, fill area between regression line and bands with a soft gradient
SHOW_DASHBOARD = True        # If True, print summary stats of each regression segment to console

# Data Download & Preprocessing
# Download daily price data and clean it up.
df = yf.download(TICKER, start=START_DATE, end=END_DATE, interval="1d")
if df.empty:
    print("No data returned for the given ticker/date range.")
    exit(1)

# Flatten MultiIndex columns if present
if isinstance(df.columns, pd.MultiIndex):
    df.columns = df.columns.map("_".join)

# Check if the SOURCE_COL exists; if not, try to find a column starting with it.
if SOURCE_COL not in df.columns:
    found = False
    for col in df.columns:
        if col.startswith(SOURCE_COL):
            df["Close"] = df[col]
            found = True
            print(f"Column '{col}' found. Renaming it to 'Close'.")
            break
    if not found:
        print(f"Column '{SOURCE_COL}' not found. Available columns: {df.columns}")
        exit(1)
else:
    df["Close"] = df[SOURCE_COL]

df.dropna(inplace=True)
if df.empty:
    print("Data is empty after cleaning. Exiting.")
    exit(1)

# Standardize columns for ease of use
df["Price"] = df["Close"]
df["Time"]  = df.index
# Convert date index to a NumPy array for positional indexing later.
dates = df.index.to_numpy()

# Build Regression Segments
# In this step, we segment the price series and fit linear regressions to each segment.
# Two modes exist:
#   - "Static": Compute regression for fixed segments (either from the first bar or periodic segments).
#   - "Rolling": Continuously update regression from an anchor point that resets periodically.

# Define helper function: get_timeframe_key, to segment the data by a chosen resolution.
def get_timeframe_key(timestamp, res="W"):
    dt = timestamp.to_pydatetime()
    if res.upper() == "W":
        y, w, _ = dt.isocalendar()
        return f"{y}-W{w}"
    elif res.upper() == "M":
        return f"{dt.year}-{dt.month:02d}"
    elif res.upper() == "D":
        return f"{dt.year}-{dt.month:02d}-{dt.day:02d}"
    elif res.upper().startswith("H"):
        hours = int(res[1:]) if len(res) > 1 else 1
        total_hours = dt.day * 24 + dt.hour
        chunk = total_hours // hours
        return f"{dt.year}-{dt.month:02d}-{dt.day:02d}-H{chunk}"
    else:
        y, w, _ = dt.isocalendar()
        return f"{y}-W{w}"

# Define helper function: linear_regression, which computes slope, intercept, correlation, and residual stdev.
def linear_regression(xs, ys):
    x = np.array(xs, dtype=float)
    y = np.array(ys, dtype=float)
    if len(x) < 2:
        return np.nan, np.nan, 0.0, 0.0
    N = len(x)
    Ex = x.sum()
    Ey = y.sum()
    Ex2 = (x * x).sum()
    Exy = (x * y).sum()
    denom = (N * Ex2 - Ex**2)
    m = (N * Exy - Ex * Ey) / denom if abs(denom) > 1e-12 else 0.0
    b = (Ey - m * Ex) / N
    r = np.corrcoef(x, y)[0, 1] if len(x) >= 2 else 0.0
    yhat = m * x + b
    residuals = y - yhat
    rss = np.mean(residuals**2)
    stdev_resid = np.sqrt(rss)
    return m, b, r, stdev_resid

# Define helper function: color_from_slope_or_corr. Returns a color based on slope and settings.
def color_from_slope_or_corr(m, r, bicolor, colA, colB, mono):
    if bicolor:
        return colA if m >= 0 else colB
    else:
        return mono

prices = df["Price"].values
xs = np.arange(len(df))  # Using bar indices for regression

segments = []  # To hold regression segments

if METHOD == "Static":
    if ANCHOR_TYPE == "First Bar":
        # Use the entire series as one segment.
        m, b, r, resid_std = linear_regression(xs, prices)
        segments.append({
            "start_idx": 0,
            "end_idx": len(df) - 1,
            "slope": m,
            "intercept": b,
            "stdev": resid_std,
            "corr": r
        })
    else:
        # Periodic segmentation based on a timeframe key (e.g., weekly).
        last_key = get_timeframe_key(df["Time"].iloc[0], ANCHOR_RES)
        seg_start = 0
        for i in range(1, len(df)):
            this_key = get_timeframe_key(df["Time"].iloc[i], ANCHOR_RES)
            if this_key != last_key or i == len(df) - 1:
                end_i = i - 1 if this_key != last_key else i
                if end_i >= seg_start:
                    x_slice = xs[seg_start:end_i + 1]
                    y_slice = prices[seg_start:end_i + 1]
                    m, b, r, resid_std = linear_regression(x_slice, y_slice)
                    segments.append({
                        "start_idx": seg_start,
                        "end_idx": end_i,
                        "slope": m,
                        "intercept": b,
                        "stdev": resid_std,
                        "corr": r
                    })
                seg_start = i
                last_key = this_key
else:
    # Rolling regression: continuously update regression from an anchor that resets periodically.
    mid_line = np.full(len(df), np.nan)
    upper_line = np.full(len(df), np.nan)
    lower_line = np.full(len(df), np.nan)
    color_line = [MONO_COLOR] * len(df)
    if ANCHOR_TYPE == "First Bar":
        seg_start = 0
        last_key = None
    else:
        last_key = get_timeframe_key(df["Time"].iloc[0], ANCHOR_RES)
        seg_start = 0
    for i in range(len(df)):
        if i == 0:
            continue
        if ANCHOR_TYPE == "Periodic":
            this_key = get_timeframe_key(df["Time"].iloc[i], ANCHOR_RES)
            if this_key != last_key:
                seg_start = i
                last_key = this_key
        x_slice = xs[seg_start:i + 1]
        y_slice = prices[seg_start:i + 1]
        m, b, r, resid_std = linear_regression(x_slice, y_slice)
        mid_val = m * i + b
        dist = resid_std * MULT
        up_val = mid_val + dist
        dn_val = mid_val - dist
        mid_line[i] = mid_val
        upper_line[i] = up_val
        lower_line[i] = dn_val
        c = COL_A if (BICOLOR and m >= 0) else (COL_B if BICOLOR else MONO_COLOR)
        color_line[i] = c
    segments.append({
        "rolling_mid": mid_line,
        "rolling_up": upper_line,
        "rolling_dn": lower_line,
        "rolling_col": color_line
    })

# Plotting
fig, ax = plt.subplots(figsize=(16, 8))
ax.set_facecolor("#1F1B1B")
ax.grid(True, alpha=0.2, color="white")
ax.set_title(f"Periodic Linear Regressions [{METHOD}] - {TICKER}", color="white")

# Use the date index as our x-axis; convert it to a NumPy array for positional indexing.
dates = df.index.to_numpy()
price = df["Price"]

ax.plot(dates, price, color="white", linewidth=1.2, label="Price")

if METHOD == "Static":
    for seg in segments:
        start_i = seg["start_idx"]
        end_i = seg["end_idx"]
        m = seg["slope"]
        b = seg["intercept"]
        stdev = seg["stdev"]
        r = seg["corr"]
        seg_x = np.arange(start_i, end_i + 1)
        seg_dates = dates[start_i:end_i + 1]
        seg_mid = m * seg_x + b
        dist = stdev * MULT
        seg_up = seg_mid + dist
        seg_down = seg_mid - dist
        c = COL_A if (BICOLOR and m >= 0) else (COL_B if BICOLOR else MONO_COLOR)
        ax.plot(seg_dates, seg_mid, color=c, linewidth=2)
        if SHOW_BANDS:
            ax.plot(seg_dates, seg_up, color=c, linestyle="--", linewidth=1)
            ax.plot(seg_dates, seg_down, color=c, linestyle="--", linewidth=1)
            if GRADIENT:
                ax.fill_between(seg_dates, seg_mid, seg_up, color=c, alpha=0.1)
                ax.fill_between(seg_dates, seg_down, seg_mid, color=c, alpha=0.1)
else:
    seg = segments[0]
    mid_line = seg["rolling_mid"]
    upper_line = seg["rolling_up"]
    lower_line = seg["rolling_dn"]
    color_line = seg["rolling_col"]
    for i in range(1, len(df)):
        c = color_line[i]
        ax.plot([dates[i - 1], dates[i]],
                [mid_line[i - 1], mid_line[i]],
                color=c, linewidth=3)
        if SHOW_BANDS:
            ax.plot([dates[i - 1], dates[i]],
                    [upper_line[i - 1], upper_line[i]],
                    color=c, linestyle="--", linewidth=1)
            ax.plot([dates[i - 1], dates[i]],
                    [lower_line[i - 1], lower_line[i]],
                    color=c, linestyle="--", linewidth=1)
            if not np.isnan(mid_line[i]) and not np.isnan(upper_line[i]):
                ax.fill_between([dates[i - 1], dates[i]],
                                [mid_line[i - 1], mid_line[i]],
                                [upper_line[i - 1], upper_line[i]],
                                color=c, alpha=0.05 if GRADIENT else 0.0)
                ax.fill_between([dates[i - 1], dates[i]],
                                [lower_line[i - 1], lower_line[i]],
                                [mid_line[i - 1], mid_line[i]],
                                color=c, alpha=0.05 if GRADIENT else 0.0)

ax.legend(loc="best", facecolor="#1F1B1B", edgecolor="white")
fig.autofmt_xdate()
plt.tight_layout()

# Optional: Print dashboard info to console.
if SHOW_DASHBOARD:
    print("=== Dashboard ===")
    print(f"Method: {METHOD}")
    print(f"Anchor Type: {ANCHOR_TYPE}")
    print(f"Anchor Period: {ANCHOR_RES}")
    print(f"Bicolor: {BICOLOR}, Gradient: {GRADIENT}")
    if METHOD == "Static":
        print(f"Segments: {len(segments)} total.")
        for i, seg in enumerate(segments):
            slope = seg["slope"]
            corr = seg["corr"]
            print(f" Segment {i}: slope = {slope:.4f}, corr = {corr:.4f}")
    else:
        print("Rolling mode => 1 big rolling segment.")

plt.show()