import numpy as np
import pandas as pd
import yfinance as yf
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.patches as mpatches
from hmmlearn.hmm import GaussianHMM
from scipy.stats import gaussian_kde

# ── PARAMETERS
TICKER   = 'TSLA'
START    = '2022-01-01'
END      = '2024-12-31'
INTERVAL = '1d'

PIVOT_K  = 5    # Pivot window size: higher = stronger, fewer pivots; lower = more, noisier pivots.
N_STATE  = 3    # Number of HMM regimes: more states = finer volatility splits, but risk of overfitting.
BAND_P   = [0.50]  # Probability bands: lower = common moves, higher = rare outliers (add more for extra bands).
GRID_NX  = 60   # KDE X grid (amplitude): higher = smoother arcs, slower plot; lower = faster, rougher arcs.
GRID_NY  = 60   # KDE Y grid (duration): same as above, for swing duration.

def download_ohlc(ticker, start, end, interval):
    df = yf.download(ticker, start=start, end=end,
                     interval=interval, auto_adjust=True, progress=False)
    # flatten MultiIndex: keep level-0 (“Open,High,…”)
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)
    df.columns = df.columns.map(str.title)
    df.dropna(subset=['Open','High','Low','Close'], inplace=True)
    return df

def detect_pivots(df, k):
    pivots, last = [], 0
    H, L = df['High'].values, df['Low'].values
    for i in range(k, len(df)-k):
        win_h = H[i-k:i+k+1]
        win_l = L[i-k:i+k+1]
        if L[i]==win_l.min() and last!= -1:
            pivots.append((i, L[i], -1)); last=-1
        elif H[i]==win_h.max() and last!=  1:
            pivots.append((i, H[i],  1)); last=1
    return pivots

def label_regimes(df, n_state):
    lr = np.log(df['Close']).diff().dropna().values.reshape(-1,1)
    model = GaussianHMM(n_components=n_state,
                        covariance_type='diag',
                        n_iter=200, random_state=0)
    model.fit(lr)
    states = model.predict(lr)
    states = np.insert(states, 0, states[0])
    return pd.Series(states, index=df.index)

def swing_stats(pivots):
    amp, dur = [], []
    for (i1,p1,t1),(i2,p2,t2) in zip(pivots, pivots[1:]):
        if t1==-1 and t2==1:
            amp.append((p2-p1)/p1); dur.append(i2-i1)
        elif t1==1 and t2==-1:
            amp.append((p1-p2)/p1); dur.append(i2-i1)
    return np.array(amp), np.array(dur)

def kde_grid(samples, levels, nx, ny):
    amp, dur = samples
    if len(amp)<15: return None
    x_max = np.percentile(amp,99)*1.2
    y_max = np.percentile(dur,99)*1.2
    xs = np.linspace(0, x_max, nx)
    ys = np.linspace(1, y_max, ny)
    X,Y = np.meshgrid(xs, ys)
    kde = gaussian_kde(np.vstack([amp,dur]))
    Z  = kde(np.vstack([X.ravel(),Y.ravel()])).reshape(X.shape)
    dx,dy = xs[1]-xs[0], ys[1]-ys[0]
    grid = np.sort(Z.ravel())[::-1]
    cdf  = np.cumsum(grid * dx * dy)
    thresh = []
    for p in levels:
        idx = np.searchsorted(cdf, p)
        idx = min(idx, len(grid)-1)
        thresh.append(grid[idx])
    return xs, ys, Z, thresh

def build_kde_by_regime(df, pivots, regimes):
    out = {}
    for st in regimes.unique():
        idxs = set(regimes[regimes==st].index)
        pv  = [p for p in pivots if df.index[p[0]] in idxs]
        out[st] = kde_grid(swing_stats(pv), BAND_P, GRID_NX, GRID_NY)
    return out  

df       = download_ohlc(TICKER, START, END, INTERVAL)
pivots   = detect_pivots(df, PIVOT_K)
regimes  = label_regimes(df, N_STATE)
kde_dict = build_kde_by_regime(df, pivots, regimes)

plt.style.use('dark_background')
fig, (ax1, ax2) = plt.subplots(
    2,1, sharex=True,
    gridspec_kw={'height_ratios':[5,1]},
    figsize=(14,8)
)

green, red = '#089981','#F23645'
width = 0.6

# ── Top: candles + contours
for dt,row in df.iterrows():
    x = mdates.date2num(dt)
    o,h,l,c = row[['Open','High','Low','Close']]
    col = green if c>=o else red
    ax1.vlines(x, l, h, color=col)
    y0 = min(o,c)
    r = mpatches.Rectangle((x-width/2,y0),
                            width, abs(c-o),
                            facecolor=col, edgecolor=col)
    ax1.add_patch(r)

for i, price, typ in pivots:
    dt   = df.index[i]
    reg  = regimes.loc[dt]
    grid = kde_dict.get(reg)
    if grid is None: continue
    xs, ys, Z, th = grid
    A, T = np.meshgrid(xs, ys)
    sign = 1 if typ==-1 else -1
    base = mdates.date2num(dt)
    Xab = base + T
    Yab = price + sign * A * price
    col = green if sign==1 else red

    for lvl, p in zip(th, BAND_P):
        cs = ax1.contour(
            Xab, Yab, Z,
            levels=[lvl],
            colors=col,
            alpha=0.4
        )
        ax1.clabel(
            cs,
            fmt={lvl: f'{int(p*100)}%'},
            inline=True,
            fontsize=8
        )

ax1.set_title(f'{TICKER} reversal probability zones')
ax1.set_ylabel('Price')

# ── Bottom: regime scatter
scatter = ax2.scatter(df.index, regimes,
                      c=regimes, cmap='tab10', s=4)
states = sorted(regimes.unique())
ax2.set_yticks(states)
ax2.set_yticklabels([str(int(s)) for s in states])
ax2.set_ylabel('Regime')
ax2.xaxis_date()
plt.subplots_adjust(bottom=0.15)
fig.autofmt_xdate()
plt.tight_layout()
plt.show()   

# ── Last closed arc stats & active pivot interpretation
if len(pivots) < 2:
    print("Not enough pivots to compute stats.")

# last closed arc = second-last → last pivot
(i1,p1,t1), (i2,p2,t2) = pivots[-2], pivots[-1]
amp_arr, dur_arr = swing_stats(pivots)
amp     = abs(p2-p1)/p1
dur     = i2 - i1
pct75_a = np.percentile(amp_arr, 75)
pct75_d = np.percentile(dur_arr, 75)
inside  = (amp <= pct75_a) and (dur <= pct75_d)

# active pivot
act_i, act_p, act_t = pivots[-1]
dt_act = df.index[act_i]
pivot_type = "Low" if act_t==-1 else "High"
regime     = regimes.iloc[act_i]

print("\n=== Last Closed Arc Summary ===")
print(f" • Start: index {i1} at price {p1:.2f}")
print(f" • End  : index {i2} at price {p2:.2f}")
print(f" • Amplitude: {amp:.2%} (vs 75th percentile: {pct75_a:.2%})")
print(f" • Duration : {dur} bars (vs 75th percentile: {int(pct75_d)} bars)")

if inside:
    print(" → This swing landed within the 75% envelope.")
    print("    → That means the move was statistically normal — both in size and duration.")
    print("    → No outlier behavior; it followed historical patterns.")
else:
    print(" → This swing moved outside the 75% envelope.")
    if amp > pct75_a and dur > pct75_d:
        print("    → The move was both larger and slower than most historical swings.")
    elif amp > pct75_a:
        print("    → Price moved farther than usual — a large extension.")
    elif dur > pct75_d:
        print("    → The move took longer than normal to complete — a slow swing.")
    print("    → This might indicate a trend continuation or breakout behavior.")

print("\n=== Active Pivot Context ===")
print(f" • Latest pivot type: {pivot_type.upper()} (printed on {dt_act.date()}, price = {act_p:.2f})")
print(f" • Market regime assigned: {regime}")

print(" • Interpretation of current regime:")
if regime == 0:
    print("    → Regime 0: Low volatility environment.")
    print("       → Expect tighter reversal zones, shallow swings, lower momentum.")
    print("       → Price is likely to respect statistical envelopes closely.")
elif regime == 1:
    print("    → Regime 1: Medium volatility.")
    print("       → Swings behave near historical averages — typical amplitudes and durations.")
    print("       → Zones offer balanced risk/reward.")
elif regime == 2:
    print("    → Regime 2: High volatility environment.")
    print("       → Expect wide zones, strong and fast moves.")
    print("       → Price may overshoot typical reversal areas.")
else:
    print("    → Unknown regime label — check the HMM configuration or input data.")

print("\nThe active arc (reversal zone) now starts from this pivot.")
print("It represents where price is statistically likely to move next — based on historical swings from similar pivots, under this same regime.")
print("Watch how price behaves relative to this arc — if it enters and reverses, it's behaving normally.")
print("If it breaks through the arc boundaries (especially the 90% zone), you're likely in a strong trend or breakout phase.")

=== Last Closed Arc Summary ===
 • Start: index 730 at price 326.59
 • End  : index 744 at price 488.54
 • Amplitude: 49.59% (vs 75th percentile: 24.77%)
 • Duration : 14 bars (vs 75th percentile: 14 bars)
 → This swing moved outside the 75% envelope.
    → Price moved farther than usual — a large extension.
    → This might indicate a trend continuation or breakout behavior.

=== Active Pivot Context ===
 • Latest pivot type: HIGH (printed on 2024-12-18, price = 488.54)
 • Market regime assigned: 2
 • Interpretation of current regime:
    → Regime 2: High volatility environment.
       → Expect wide zones, strong and fast moves.
       → Price may overshoot typical reversal areas.

The active arc (reversal zone) now starts from this pivot.
It represents where price is statistically likely to move next — based on historical swings from similar pivots, under this same regime.
Watch how price behaves relative to this arc — if it enters and reverses, it's behaving normally.
If it breaks through the arc boundaries (especially the 90% zone), you're likely in a strong trend or breakout phase.