"""
Метод Монте-Карло: от pi до интегралов
Monte Carlo estimation of pi + complex integral, convergence visualization
1080x1080, 25 seconds, 30fps, dark theme
"""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter
from matplotlib.patches import Circle, FancyBboxPatch
from matplotlib.collections import PathCollection

# === Config ===
FPS = 30
DURATION = 25
N_FRAMES = FPS * DURATION
W, H = 1080, 1080
DPI = 120

BG_COLOR = '#1a1a2e'
TEXT_COLOR = '#e0e0e0'
INSIDE_COLOR = '#00d2ff'
OUTSIDE_COLOR = '#ff6b6b'
CONVERGENCE_COLOR = '#ffd93d'
INTEGRAL_COLOR = '#6bcb77'
TRUE_PI_COLOR = '#ff9ff3'

# === Pre-compute Monte Carlo for pi ===
np.random.seed(42)
TOTAL_POINTS = 8000  # total points to throw

# Generate all points at once
all_x = np.random.uniform(-1, 1, TOTAL_POINTS)
all_y = np.random.uniform(-1, 1, TOTAL_POINTS)
all_inside = (all_x**2 + all_y**2) <= 1.0

# Compute running pi estimate
cumsum_inside = np.cumsum(all_inside)
point_indices = np.arange(1, TOTAL_POINTS + 1)
pi_estimates = 4.0 * cumsum_inside / point_indices

# === Pre-compute Monte Carlo for integral ===
# Integral of f(x) = sin(x) * exp(-x^2/2) from -3 to 3
# True value computed numerically
from scipy import integrate
def integrand(x):
    return np.sin(x) * np.exp(-x**2 / 2)

true_integral, _ = integrate.quad(integrand, -3, 3)
print(f"True integral value: {true_integral:.6f}")
# This is ~0 by symmetry since sin is odd and exp(-x^2/2) is even
# Let's use a more interesting asymmetric function
def integrand2(x):
    return np.sin(x)**2 * np.exp(-x**2 / 4) + 0.3

true_integral2, _ = integrate.quad(integrand2, -3, 3)
print(f"True integral value (f2): {true_integral2:.6f}")

TOTAL_POINTS_INT = 6000
int_x = np.random.uniform(-3, 3, TOTAL_POINTS_INT)
int_y_max = 1.6  # upper bound for rejection sampling visualization
int_y = np.random.uniform(0, int_y_max, TOTAL_POINTS_INT)

int_fx = integrand2(int_x)
int_inside = int_y <= int_fx

# Running estimate: area = (x_range * y_max) * (n_inside / n_total)
x_range = 6.0  # from -3 to 3
cumsum_int_inside = np.cumsum(int_inside)
int_estimates = x_range * int_y_max * cumsum_int_inside / np.arange(1, TOTAL_POINTS_INT + 1)

# === Map frames to points ===
# Phase 1 (0-60%): pi estimation
# Phase 2 (60-100%): integral estimation
PHASE1_END = int(N_FRAMES * 0.55)
PHASE2_START = int(N_FRAMES * 0.55)
TRANSITION_FRAMES = int(N_FRAMES * 0.05)

def frame_to_pi_points(frame):
    """How many pi points to show at this frame"""
    if frame >= PHASE1_END:
        return TOTAL_POINTS
    progress = frame / PHASE1_END
    # Accelerating: start slow, speed up
    return max(1, int(TOTAL_POINTS * (progress ** 1.5)))

def frame_to_int_points(frame):
    """How many integral points to show at this frame"""
    if frame < PHASE2_START:
        return 0
    progress = (frame - PHASE2_START) / (N_FRAMES - PHASE2_START)
    return max(1, int(TOTAL_POINTS_INT * (progress ** 1.3)))

# === Create figure ===
fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
fig.patch.set_facecolor(BG_COLOR)

gs = fig.add_gridspec(2, 2, hspace=0.28, wspace=0.25,
                      left=0.08, right=0.95, top=0.92, bottom=0.06)

ax_pi = fig.add_subplot(gs[0, 0])        # pi scatter
ax_pi_conv = fig.add_subplot(gs[0, 1])   # pi convergence
ax_int = fig.add_subplot(gs[1, 0])       # integral scatter
ax_int_conv = fig.add_subplot(gs[1, 1])  # integral convergence

fig.suptitle('Метод Монте-Карло: от числа π до интегралов',
             fontsize=16, color=TEXT_COLOR, fontweight='bold', y=0.97)

# === Pi scatter setup ===
ax_pi.set_facecolor(BG_COLOR)
ax_pi.set_xlim(-1.15, 1.15)
ax_pi.set_ylim(-1.15, 1.15)
ax_pi.set_aspect('equal')
ax_pi.set_title('Оценка числа π', fontsize=12, color=INSIDE_COLOR, fontweight='bold', pad=6)
ax_pi.tick_params(colors='#555555', labelsize=7)
for spine in ax_pi.spines.values():
    spine.set_color('#333355')

# Draw unit circle
theta = np.linspace(0, 2*np.pi, 200)
ax_pi.plot(np.cos(theta), np.sin(theta), color=INSIDE_COLOR, linewidth=1.5, alpha=0.5)
# Draw square
sq = plt.Rectangle((-1, -1), 2, 2, fill=False, edgecolor='#555577', linewidth=1, linestyle='--')
ax_pi.add_patch(sq)

pi_inside_scatter = ax_pi.scatter([], [], s=3, c=INSIDE_COLOR, alpha=0.5, zorder=2)
pi_outside_scatter = ax_pi.scatter([], [], s=3, c=OUTSIDE_COLOR, alpha=0.3, zorder=2)

pi_text = ax_pi.text(0.5, -0.12, '', transform=ax_pi.transAxes, ha='center',
                     fontsize=13, color=CONVERGENCE_COLOR, fontfamily='monospace',
                     fontweight='bold')

# === Pi convergence setup ===
ax_pi_conv.set_facecolor(BG_COLOR)
ax_pi_conv.set_xlim(1, TOTAL_POINTS)
ax_pi_conv.set_ylim(2.5, 3.8)
ax_pi_conv.set_xscale('log')
ax_pi_conv.set_title('Сходимость к π', fontsize=12, color=CONVERGENCE_COLOR,
                     fontweight='bold', pad=6)
ax_pi_conv.set_xlabel('N точек', fontsize=9, color=TEXT_COLOR)
ax_pi_conv.tick_params(colors='#555555', labelsize=7)
for spine in ax_pi_conv.spines.values():
    spine.set_color('#333355')

# True pi line
ax_pi_conv.axhline(y=np.pi, color=TRUE_PI_COLOR, linewidth=1.5, linestyle='--', alpha=0.7)
ax_pi_conv.text(0.02, 0.87, f'π = {np.pi:.5f}', transform=ax_pi_conv.transAxes,
                fontsize=9, color=TRUE_PI_COLOR, fontfamily='monospace')

# 1/sqrt(n) bounds
n_arr = np.arange(1, TOTAL_POINTS + 1)
upper = np.pi + 4.0 / np.sqrt(n_arr)
lower = np.pi - 4.0 / np.sqrt(n_arr)
ax_pi_conv.fill_between(n_arr, lower, upper, alpha=0.08, color=CONVERGENCE_COLOR)
ax_pi_conv.text(0.55, 0.12, '~1/√n', transform=ax_pi_conv.transAxes,
                fontsize=11, color=CONVERGENCE_COLOR, alpha=0.5, fontfamily='monospace')

pi_conv_line, = ax_pi_conv.plot([], [], color=CONVERGENCE_COLOR, linewidth=1.5)

# === Integral scatter setup ===
ax_int.set_facecolor(BG_COLOR)
ax_int.set_xlim(-3.3, 3.3)
ax_int.set_ylim(-0.1, int_y_max + 0.1)
ax_int.set_title('Оценка интеграла', fontsize=12, color=INTEGRAL_COLOR,
                 fontweight='bold', pad=6)
ax_int.tick_params(colors='#555555', labelsize=7)
for spine in ax_int.spines.values():
    spine.set_color('#333355')

# Draw the function
x_func = np.linspace(-3, 3, 500)
y_func = integrand2(x_func)
ax_int.plot(x_func, y_func, color=INTEGRAL_COLOR, linewidth=2, alpha=0.8)
ax_int.fill_between(x_func, 0, y_func, alpha=0.1, color=INTEGRAL_COLOR)

int_inside_scatter = ax_int.scatter([], [], s=3, c=INTEGRAL_COLOR, alpha=0.5, zorder=2)
int_outside_scatter = ax_int.scatter([], [], s=3, c=OUTSIDE_COLOR, alpha=0.2, zorder=2)

int_text = ax_int.text(0.5, -0.12, '', transform=ax_int.transAxes, ha='center',
                       fontsize=11, color=INTEGRAL_COLOR, fontfamily='monospace',
                       fontweight='bold')

# === Integral convergence setup ===
ax_int_conv.set_facecolor(BG_COLOR)
ax_int_conv.set_xlim(1, TOTAL_POINTS_INT)
ax_int_conv.set_xscale('log')
ax_int_conv.set_title('Сходимость интеграла', fontsize=12, color=INTEGRAL_COLOR,
                      fontweight='bold', pad=6)
ax_int_conv.set_xlabel('N точек', fontsize=9, color=TEXT_COLOR)
ax_int_conv.tick_params(colors='#555555', labelsize=7)
for spine in ax_int_conv.spines.values():
    spine.set_color('#333355')

ax_int_conv.axhline(y=true_integral2, color=TRUE_PI_COLOR, linewidth=1.5,
                    linestyle='--', alpha=0.7)
ax_int_conv.text(0.02, 0.87, f'True = {true_integral2:.3f}', transform=ax_int_conv.transAxes,
                 fontsize=9, color=TRUE_PI_COLOR, fontfamily='monospace')
ax_int_conv.set_ylim(true_integral2 * 0.5, true_integral2 * 1.5)

int_conv_line, = ax_int_conv.plot([], [], color=INTEGRAL_COLOR, linewidth=1.5)

counter_text = fig.text(0.5, 0.015, '', ha='center', va='bottom',
                        fontsize=11, color='#888888', fontfamily='monospace')

def animate(frame):
    # === Phase 1: Pi ===
    n_pi = frame_to_pi_points(frame)

    mask_in = all_inside[:n_pi]
    mask_out = ~mask_in

    x_in = all_x[:n_pi][mask_in]
    y_in = all_y[:n_pi][mask_in]
    x_out = all_x[:n_pi][~all_inside[:n_pi]]
    y_out = all_y[:n_pi][~all_inside[:n_pi]]

    pi_inside_scatter.set_offsets(np.column_stack([x_in, y_in]) if len(x_in) > 0
                                  else np.empty((0, 2)))
    pi_outside_scatter.set_offsets(np.column_stack([x_out, y_out]) if len(x_out) > 0
                                   else np.empty((0, 2)))

    pi_est = pi_estimates[n_pi - 1]
    pi_text.set_text(f'π ≈ {pi_est:.5f}  (N={n_pi})')

    # Convergence line
    pi_conv_line.set_data(point_indices[:n_pi], pi_estimates[:n_pi])

    # === Phase 2: Integral ===
    n_int = frame_to_int_points(frame)

    if n_int > 0:
        mask_in_int = int_inside[:n_int]
        ix_in = int_x[:n_int][mask_in_int]
        iy_in = int_y[:n_int][mask_in_int]
        ix_out = int_x[:n_int][~mask_in_int]
        iy_out = int_y[:n_int][~mask_in_int]

        int_inside_scatter.set_offsets(np.column_stack([ix_in, iy_in]) if len(ix_in) > 0
                                       else np.empty((0, 2)))
        int_outside_scatter.set_offsets(np.column_stack([ix_out, iy_out]) if len(ix_out) > 0
                                        else np.empty((0, 2)))

        est = int_estimates[n_int - 1]
        int_text.set_text(f'∫f(x)dx ≈ {est:.3f}  (N={n_int})')

        int_conv_line.set_data(np.arange(1, n_int+1), int_estimates[:n_int])
    else:
        int_inside_scatter.set_offsets(np.empty((0, 2)))
        int_outside_scatter.set_offsets(np.empty((0, 2)))
        int_text.set_text('')

    total = n_pi + n_int
    counter_text.set_text(f'Всего точек: {total}')

    return [pi_inside_scatter, pi_outside_scatter, pi_text, pi_conv_line,
            int_inside_scatter, int_outside_scatter, int_text, int_conv_line,
            counter_text]

anim = FuncAnimation(fig, animate, frames=N_FRAMES, interval=1000/FPS, blit=False)

outpath = '/root/Strategy/content/generated/monte_carlo_animation.mp4'
writer = FFMpegWriter(fps=FPS, bitrate=3000,
                      extra_args=['-vcodec', 'libx264', '-pix_fmt', 'yuv420p'])
anim.save(outpath, writer=writer)
plt.close()
print(f"Saved: {outpath}")