import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

if len(sys.argv) != 2:
    print('Usage: python moves-to-go.py <filename>')
    sys.exit(1)

csv_file = sys.argv[1]
df = pd.read_csv(csv_file)

# Ensure correct column order and sorting.
df = df.sort_values('Evaluation')

# Fill gaps by linear interpolation first.
df['MovesToGoMedian'] = df['MovesToGoMedian'].interpolate('linear')

# Apply moving-average smoothing. All the little peaks in the curve are
# probaly not explainable and we could try to smooth them out further. But
# having a little bit of randomness in the time allocation is not necessarily
# bad, and so we don't exaggerate.
window = 25  # Tune this: larger = smoother, smaller = more detail
df['Smoothed'] = df['MovesToGoMedian'].rolling(window=window, center=True).mean()
df['Smoothed'] = df['Smoothed'].bfill().ffill()

# Compute statistics.
r2 = np.corrcoef(df['Evaluation'], df['MovesToGoMedian'])[0, 1] ** 2
r2_smoothed = np.corrcoef(df['Evaluation'], df['Smoothed'])[0, 1] ** 2

corr_original = np.corrcoef(df['MovesToGoMedian'], df['Evaluation'])[0, 1]
corr_smoothed = np.corrcoef(df['Smoothed'], df['Evaluation'])[0, 1]

#print(f"R² (original): {r2:.4f}")
#print(f"R² (smoothed): {r2_smoothed:.4f}")
#print(f"Correlation (original): {corr_original:.4f}")
#print(f"Correlation (smoothed): {corr_smoothed:.4f}")

# Plot results
plt.figure(figsize=(8, 6))
plt.plot(df['Evaluation'], df['MovesToGoMedian'], 'gray', alpha=0.4, label='Original / interpolated')
plt.plot(df['Evaluation'], df['Smoothed'], 'red', label=f'Moving average (window={window})')
plt.xlabel('Evaluation (centipawns)')
plt.ylabel('MovesToGoMedian')
plt.title('Smoothed Moves-to-Go vs Evaluation')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

# Output the table.
print("# This file is generated by scripts/moves-to-go.p[ly].  Do not edit!\n")

print("package Chess::Plisco::Engine::TimeControl::MovesToGo;\n")
print("use strict;\n")

print("use constant MOVES_TO_GO => [")

lines = []
line = []
for val in df['Smoothed']:
    if val < 10:
        break
    line.append(f"{val:.2f}")
    if len(line) >= 9:
        lines.append(', '.join(line))
        line = []

if line:
    lines.append(', '.join(line))

print("\t" + ",\n\t".join(lines))
print("];\n\n1;")
