"""Reproduce the synthetic example and plots in the QuantSteps guide.

Run: python heteroscedasticity.py
Dependencies: numpy, scipy, statsmodels, matplotlib.
Outputs are saved beside this file; no market or personal data is downloaded.
"""
from pathlib import Path
import csv
import json
import platform

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import scipy
import statsmodels
import statsmodels.api as sm
from statsmodels.stats.diagnostic import het_breuschpagan, het_white

OUT = Path(__file__).resolve().parent
rng = np.random.default_rng(20261003)
n = 600
x = np.linspace(0.5, 5.0, n)
X = sm.add_constant(x)
noise = rng.standard_normal(n)
sigma = 0.3 + 0.65 * x
y = 2.0 + 0.8 * x + sigma * noise
ols = sm.OLS(y, X).fit()
hc3 = ols.get_robustcov_results(cov_type='HC3', use_t=True)
# Oracle weights: known because WE generated sigma, not because data reveal it.
wls = sm.WLS(y, X, weights=1.0 / sigma**2).fit()
bp = het_breuschpagan(ols.resid, X, robust=True)
white = het_white(ols.resid, X)
assert np.allclose(ols.params, hc3.params)
assert np.allclose(ols.fittedvalues, hc3.fittedvalues)

results = {
    'seed': 20261003, 'n': n,
    'versions': {'python': platform.python_version(), 'numpy': np.__version__,
                 'scipy': scipy.__version__, 'statsmodels': statsmodels.__version__,
                 'matplotlib': matplotlib.__version__},
    'models': {name: {'coef': model.params.tolist(), 'se': model.bse.tolist(),
                     'ci95': model.conf_int().tolist()}
               for name, model in [('OLS', ols), ('OLS_HC3_t', hc3), ('WLS_oracle', wls)]},
    'tests': {'Koenker_BP': dict(zip(['LM', 'LM_p', 'F', 'F_p'], map(float, bp))),
              'White': dict(zip(['LM', 'LM_p', 'F', 'F_p'], map(float, white)))},
}
(OUT / 'heteroscedasticity-results.json').write_text(json.dumps(results, indent=2), encoding='utf-8')
with (OUT / 'heteroscedasticity-data.csv').open('w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(['x', 'y', 'known_sigma', 'OLS_fitted', 'OLS_residual'])
    writer.writerows(zip(x, y, sigma, ols.fittedvalues, ols.resid))

plt.rcParams.update({'font.family': 'DejaVu Sans', 'font.size': 10,
                     'figure.facecolor': '#151515', 'axes.facecolor': '#151515',
                     'text.color': '#e5e7eb', 'axes.labelcolor': '#bdbdbd',
                     'xtick.color': '#bdbdbd', 'ytick.color': '#bdbdbd',
                     'axes.edgecolor': '#555555', 'savefig.facecolor': '#151515'})
fig, axes = plt.subplots(2, 2, figsize=(11, 7), sharex=True, sharey=True)
patterns = [('Constant variance', np.full(n, 1.8)),
            ('Widening variance (main example)', sigma),
            ('Narrowing variance', 3.6 - 0.65 * x),
            ('Group variance', np.where(x < 2.75, 0.7, 2.8))]
for ax, (title, sd) in zip(axes.flat, patterns):
    sample = 2 + 0.8 * x + sd * noise
    model = sm.OLS(sample, X).fit()
    ax.scatter(x, model.resid, s=9, color='#93c5fd', alpha=0.58, edgecolors='none')
    ax.axhline(0, color='#a3a3a3', lw=0.8, ls='--')
    ax.set_title(title, loc='left', fontsize=12, pad=10)
    ax.set_xlabel('Predictor x')
    ax.set_ylabel('OLS residual')
    ax.spines[['top', 'right']].set_visible(False)
fig.suptitle('Same conditional mean; different error variances', fontsize=16, x=0.05, ha='left')
fig.tight_layout(rect=[0, 0, 1, 0.94])
fig.savefig(OUT / 'heteroscedasticity-residuals.png', dpi=160)
plt.close(fig)
print(json.dumps(results, indent=2))