import numpy as np
import yfinance as yf
import torch
import math
import matplotlib.pyplot as plt

# ============================================================
# CCT-GRADIENT BTC FORECASTING
# Analogy: BTC price = "semiprime" with hidden factors
# Factorization: c = p × q  →  BTC: P = B × S
# B = "Base factor" (fundamental value)
# S = "Sentiment factor" (market multiplier)
# ============================================================

def cct_btc_gradient_forecast(prices, days_ahead=40, lr=0.05, num_restarts=5):
    """
    BTC forecasting via gradient descent on the price gap.
    
    Key Insight:
    - Treat BTC price P as having a "hidden factorization" structure
    - Find the "gap" d between price and market equilibrium
    - The loss L(d) = sin²(π·s) where s = √(d² + 4P)
    - Gradient descent collapses d to find the true price manifold
    """
    
    P = prices[-1]  # Current price
    
    # Loss function: identical to factorization but with price instead of c
    # s represents the "market integer invariant"
    def loss_fn(d, P_val):
        """
        d = gap between price and equilibrium
        P_val = current BTC price
        s = √(d² + 4P) is the "market invariant" that should be integer-like
        """
        s = torch.sqrt(d**2 + 4.0 * P_val)
        # sin²(π·s) is 0 when s is integer
        # We want s to be integer (stable market state)
        loss = torch.sin(torch.pi * s) ** 2
        return loss, s
    
    P_t = torch.tensor(float(P), requires_grad=False)
    
    best_d = None
    best_loss = float('inf')
    all_trajectories = []
    
    for restart in range(num_restarts):
        d = torch.randn(1) * np.sqrt(P) / 10  # Scale init by √P
        d = d.detach().requires_grad_()
        optimizer = torch.optim.Adam([d], lr=lr)
        
        trajectory = []
        
        for step in range(500):
            optimizer.zero_grad()
            loss, s = loss_fn(d, P_t)
            loss.backward()
            optimizer.step()
            
            with torch.no_grad():
                d.clamp_(min=0.0)
            
            trajectory.append({
                'step': step,
                'd': d.item(),
                'loss': loss.item(),
                's': s.item()
            })
            
            if loss.item() < best_loss:
                best_loss = loss.item()
                best_d = d.item()
        
        all_trajectories.append(trajectory)
        
        if best_loss < 1e-10:
            break
    
    # Recover "factors" of price (analogous to p, q)
    d_star = round(best_d)
    s_sq = d_star**2 + 4 * P
    s = int(round(math.sqrt(s_sq)))
    
    # Analogous to p = (s - d)/2, q = (s + d)/2
    base = (s - d_star) // 2   # "Fundamental value"
    sentiment = (s + d_star) // 2  # "Market sentiment"
    
    return {
        'd_star': d_star,
        's': s,
        'base_factor': base,
        'sentiment_factor': sentiment,
        'loss': best_loss,
        'trajectories': all_trajectories
    }


def predict_multi_step(prices, days_ahead=40):
    """
    CCT-ODE Multi-Step Forecast:
    Each future day is predicted by gradient descent on the current state.
    This creates a "chain of collapses" - each day collapses independently.
    """
    forecasts = []
    current_prices = prices.copy()
    
    for day in range(days_ahead):
        # Get current price
        P = current_prices[-1]
        
        # Run gradient descent to find equilibrium
        result = cct_btc_gradient_forecast(current_prices, days_ahead=1, num_restarts=3)
        
        # The "base factor" is our predicted equilibrium price
        predicted_price = result['base_factor']
        
        # Add small perturbation (mimics market noise)
        predicted_price *= (1 + 0.02 * np.random.randn())
        
        forecasts.append(predicted_price)
        current_prices = np.append(current_prices, predicted_price)
        
        print(f"Day {day+1}: P={P:.2f} → Predicted={predicted_price:.2f} | "
              f"Base Factor={result['base_factor']:.2f}, Gap d={result['d_star']}")
    
    return np.array(forecasts)


# ============================================================
# LOAD BTC DATA
# ============================================================
print("Downloading BTC data...")
btc = yf.download("BTC-USD", period="6mo", interval="1d")
close = btc['Close'].dropna().values.flatten()
print(f"Loaded {len(close)} days of BTC data")
print(f"Current price: ${close[-1]:,.2f}")

# ============================================================
# RUN CCT-GRADIENT FORECAST
# ============================================================
print("\n" + "="*60)
print("CCT-GRADIENT BTC FORECAST")
print("="*60)

# Single-step analysis
result = cct_btc_gradient_forecast(close, days_ahead=1, num_restarts=10)

print(f"\nCCT Analysis of Current Price:")
print(f"  Current Price P: ${close[-1]:,.2f}")
print(f"  Optimal Gap d*: {result['d_star']}")
print(f"  Market Invariant s: {result['s']}")
print(f"  Base Factor (Fundamental): ${result['base_factor']:,.2f}")
print(f"  Sentiment Factor: {result['sentiment_factor']:.2f}")
print(f"  Loss (Collapsed): {result['loss']:.2e}")

# Multi-step forecast
print(f"\nGenerating 40-day forecast...")
forecast = predict_multi_step(close, days_ahead=40)

# ============================================================
# VISUALIZATION
# ============================================================
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# 1. Gradient Descent Trajectories
ax = axes[0, 0]
for i, traj in enumerate(result['trajectories']):
    losses = [t['loss'] for t in traj]
    ax.plot(losses, alpha=0.7, label=f'Restart {i+1}')
ax.set_yscale('log')
ax.set_title("CCT Gradient Descent: Loss Convergence")
ax.set_xlabel("Step")
ax.set_ylabel("Loss (sin²)")
ax.legend()

# 2. Gap d over iterations
ax = axes[0, 1]
for i, traj in enumerate(result['trajectories']):
    ds = [t['d'] for t in traj]
    ax.plot(ds, alpha=0.7, label=f'Restart {i+1}')
ax.axhline(y=result['d_star'], color='red', linestyle='--', label=f'Final d*={result["d_star"]}')
ax.set_title("CCT Gradient Descent: Gap d Convergence")
ax.set_xlabel("Step")
ax.set_ylabel("Gap d")
ax.legend()

# 3. Historical + Forecast
ax = axes[1, 0]
days_hist = np.arange(len(close))
days_fc = np.arange(len(close), len(close) + 40)

ax.plot(days_hist, close, 'b-', label='Historical', linewidth=1)
ax.plot(days_fc, forecast, 'r-', label='CCT Forecast', linewidth=2)
ax.fill_between(days_fc, forecast * 0.9, forecast * 1.1, alpha=0.2, color='red')
ax.axvline(x=len(close)-1, color='gray', linestyle=':', label='Today')
ax.set_title("BTC-USD: Historical + 40-Day CCT Forecast")
ax.set_xlabel("Day")
ax.set_ylabel("Price (USD)")
ax.legend()

# 4. Structure: Factorization Analogy
ax = axes[1, 1]
ax.axis('off')
text = f"""
CCT-GRADIENT STRUCTURE
═══════════════════════════════════════════

BTC Price P = Base Factor × Sentiment Factor

Current Price: ${close[-1]:,.2f}

┌─────────────────────────────────────────────┐
│  Loss: L(d) = sin²(π·s)                     │
│  where s = √(d² + 4P)                       │
│                                             │
│  Gradient: ∂L/∂d = π·sin(2πs)·d/s           │
│                                             │
│  Converges when s → integer (stable state)  │
└─────────────────────────────────────────────┘

Analogy to Factorization:
• Factorization: c = p × q, find gap d = q - p
• BTC: P = B × S, find gap d = S - B

Both use: s² = d² + 4c  (or 4P)

The "integer constraint" forces collapse
to the true factor structure.
"""
ax.text(0.1, 0.9, text, transform=ax.transAxes, fontsize=10,
        verticalalignment='top', fontfamily='monospace',
        bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.tight_layout()
plt.show()

# ============================================================
# PRINT FORECAST TABLE
# ============================================================
print("\n" + "="*60)
print("40-Day BTC Forecast (CCT-Gradient)")
print("="*60)
print(f"{'Day':<6} {'Forecast':<20} {'Base Factor':<15}")
print("-"*60)
for d in [1, 5, 10, 20, 40]:
    if d <= len(forecast):
        print(f"{d:<6} ${forecast[d-1]:>15,.2f}")
print("="*60)
