import math
import numpy as np
import matplotlib.pyplot as plt
from numba import njit

# Setup
%matplotlib inline
%config InlineBackend.close_figures = False

# Define the IsingSimulator class
class IsingSimulator:
    """Numba-accelerated 2D Ising model Monte Carlo simulator."""

    def __init__(self, batch_size=1, system_size=16, spin_config=None, J=1.0, T=2.0, h=0.0):
        self.batch_size = batch_size
        self.system_size = system_size
        self.J = J  # coupling constant
        self.T = T  # temperature (k_B=1)
        self.h = h  # external field (Ising convention)
        self.reset(spin_config)

    @property
    def magnetization(self):
        """(mean, std) over batches. mean = batch mean of per-config magnetization; std for error bars."""
        m_per_batch = self.spin_config.astype(np.float64).mean(axis=(1, 2))
        return float(m_per_batch.mean()), float(m_per_batch.std(ddof=1)) if len(m_per_batch) > 1 else 0.0

    def reset(self, spin_config=None):
        if spin_config is None:
            self.spin_config = 2 * np.random.randint(2, size=(self.batch_size, self.system_size, self.system_size)) - 1
        else:
            if hasattr(spin_config, 'numpy'):
                spin_config = spin_config.numpy()
            self.batch_size = spin_config.shape[0]
            self.system_size = spin_config.shape[1]
            self.spin_config = np.asarray(spin_config, dtype=np.int8)

    def update(self, steps=None):
        """Single-spin Metropolis update (for step-by-step teaching demo)."""
        N, L = self.batch_size, self.system_size
        if steps is None:
            steps = L * L
        s_flat = np.ascontiguousarray(self.spin_config.reshape(N, -1), dtype=np.int8)
        self._update_kernel(s_flat, L, self.J, self.T, self.h, steps)
        self.spin_config = s_flat.reshape(N, L, L)

    def equilibrate(self, sweeps=100):
        """Checkerboard Gibbs update for fast equilibration. Use for hysteresis/magnetization curves."""
        L = self.system_size
        s = np.ascontiguousarray(self.spin_config, dtype=np.int8)
        beta = 1.0 / self.T
        mask = (np.add.outer(np.arange(L), np.arange(L)) % 2 == 0)
        masks = [mask, ~mask]  # even and odd sublattices
        for _ in range(sweeps):
            for mask in masks:
                sum_n = np.roll(s, 1, axis=1) + np.roll(s, -1, axis=1) + np.roll(s, 1, axis=2) + np.roll(s, -1, axis=2)
                H = self.J * sum_n.astype(np.float64) + self.h
                P_up = 1.0 / (1.0 + np.exp(-2.0 * beta * H))
                u = np.random.random(s.shape)
                new_spin = (2 * (u < P_up).astype(np.int8) - 1)
                s[:, mask] = new_spin[:, mask]
        self.spin_config = s

    def redraw(self, batch_idx=0):
        """Update display."""
        if not hasattr(self, '_fig') or self._fig is None:
            self._fig, self._ax = plt.subplots(figsize=(3, 3))
            self._im = self._draw_spin_config(self._fig, self._ax, batch_idx)
            self._fig.canvas.header_visible = False
            self._fig.canvas.toolbar_visible = False
            self._fig.canvas.footer_visible = False
        else:
            self._im.set_data((1 - self.spin_config[batch_idx]) / 2)
        self._fig.canvas.draw_idle()

    def _draw_spin_config(self, fig, ax, batch_idx=0):
        """Draw spin config with common styling. Returns the AxesImage for in-place updates."""
        data = (1 - self.spin_config[batch_idx]) / 2  # -1→black, +1→white
        im = ax.imshow(data, cmap='binary', vmin=0, vmax=1)
        ax.set_xticks([])
        ax.set_yticks([])
        for spine in ax.spines.values():
            spine.set_visible(True)
            spine.set_color('black')
            spine.set_linewidth(1.5)
        fig.suptitle('Spin configuration', fontsize=10)
        fig.text(0.5, 0.02, 'Black: spin ↓, White: spin ↑', ha='center', fontsize=9)
        fig.subplots_adjust(left=0, right=1, top=0.9, bottom=0.1)
        return im

    @staticmethod
    @njit
    def _update_kernel(s_flat, L, J, T, h, steps):
        """Compiled MC update. s_flat: (N, L*L) int8 spin config, modified in place."""
        N, n_sites = s_flat.shape[0], L * L
        for _ in range(steps):
            for b in range(N):
                i = np.random.randint(0, n_sites)
                i1, i2 = i // L, i % L
                il = ((i1 - 1) % L) * L + i2
                ir = ((i1 + 1) % L) * L + i2
                ib = i1 * L + (i2 - 1) % L
                it = i1 * L + (i2 + 1) % L
                s_i = s_flat[b, i]
                sum_n = s_flat[b, il] + s_flat[b, ir] + s_flat[b, ib] + s_flat[b, it]
                dE = 2.0 * J * s_i * sum_n + 2.0 * h * s_i
                if np.random.random() <= min(1.0, math.exp(-dE / T)):
                    s_flat[b, i] = -s_i

import matplotlib.pyplot as plt
from ipywidgets import Play, Button, FloatSlider, VBox, HBox, Layout, Output
from IPython.display import display, clear_output

plt.close('all')
plt.ioff()

sim = IsingSimulator(batch_size=1, system_size=64, T=2.0, h=0.0)

# Create figure once for reuse
fig, ax = plt.subplots()
fig.set_size_inches(3, 3)
data = (1 - sim.spin_config[0]) / 2  # -1→1 (white), +1→0 (black)
im = ax.imshow(data, cmap='binary', vmin=0, vmax=1)
ax.set_xticks([])
ax.set_yticks([])
for spine in ax.spines.values():
    spine.set_visible(True)
    spine.set_color('black')
    spine.set_linewidth(1.5)
fig.suptitle('Spin configuration', fontsize=10)
fig.text(0.5, 0.02, 'Black: spin ↓, White: spin ↑', ha='center', fontsize=9)
fig.subplots_adjust(left=0, right=1, top=0.9, bottom=0.1)

out = Output(layout=Layout(width='320px', min_width='320px', height='340px'))

def update_display(step, T, B):
    plt.ioff()
    sim.T = T
    sim.h = B
    if step > 0:
        sim.update()
    im.set_data((1 - sim.spin_config[0]) / 2)
    fig.set_size_inches(3, 3)
    with out:
        clear_output(wait=True)
        display(fig)
    plt.ion()

# Play widget – click its play button to animate; drives step count for update_display
play = Play(value=0, min=0, max=999999999, step=1, interval=16, show_repeat=False)
T_slider = FloatSlider(value=2.0, min=0.5, max=4.0, step=0.01, description='T:', style={'description_width': '20px'}, continuous_update=False)
B_slider = FloatSlider(value=0.0, min=-2.0, max=2.0, step=0.01, description='B:', style={'description_width': '20px'}, continuous_update=False)

def on_change(change):
    update_display(play.value, T_slider.value, B_slider.value)
play.observe(on_change, 'value')
T_slider.observe(on_change, 'value')
B_slider.observe(on_change, 'value')

# Initial display (observe only fires on change)
with out:
    clear_output(wait=True)
    display(fig)

def on_refresh(b):
    plt.ioff()
    sim.reset()
    with out:
        clear_output(wait=True)
        im.set_data((1 - sim.spin_config[0]) / 2)
        fig.set_size_inches(3, 3)
        display(fig)
    plt.ion()

refresh_btn = Button(icon='refresh', description='Refresh')
refresh_btn.on_click(on_refresh)

controls_col = VBox([play, T_slider, B_slider, refresh_btn], layout=Layout(align_items='flex-start', flex='0 0 auto'))
main_hbox = HBox([out, controls_col], layout=Layout(gap='8px', align_items='flex-start'))

# out created first so update_display/on_refresh can reference it
display(main_hbox)
# Prevent inline backend from auto-displaying fig at cell end (would leak to cell output)
plt.figure(); plt.close()
plt.ion();

import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import interact, FloatSlider
from IPython.display import display

# Hysteresis parameters
B_MIN, B_MAX = -2.0, 2.0
N_STEPS = 21
BATCH_SIZE, L = 16, 32
EQUIL_SWEEPS = 10

# Build B sequence: -2→2, 2→-2, ...
up = np.linspace(B_MIN, B_MAX, N_STEPS)
down = np.linspace(B_MAX, B_MIN, N_STEPS)
B_values = np.concatenate([up, down[1:]])

@interact(T=FloatSlider(value=2.0, min=0.5, max=4.0, step=0.01, description='T:', style={'description_width': '20px'}, continuous_update=False))
def run_hysteresis(T):
    plt.ioff()
    sim = IsingSimulator(batch_size=BATCH_SIZE, system_size=L, T=T, h=0.0)
    B_list, M_list, M_err_list = [], [], []
    for h in B_values:
        sim.h, sim.T = h, T
        sim.equilibrate(sweeps=EQUIL_SWEEPS)
        m, m_std = sim.magnetization
        B_list.append(h)
        M_list.append(m)
        M_err_list.append(m_std)
    fig, ax = plt.subplots(figsize=(5, 3))
    ax.set_xlabel('B')
    ax.set_ylabel('M')
    ax.axhline(0, color='gray', ls='--', alpha=0.5)
    ax.axvline(0, color='gray', ls='--', alpha=0.5)
    ax.errorbar(B_list, M_list, yerr=M_err_list, fmt='o-', capsize=2, markersize=4)
    ax.set_title(f'Hysteresis (T={T:.2f})')
    plt.tight_layout()
    display(fig)
    plt.close(fig)
    plt.ion()

# M vs T at fixed small B: Curie-Weiss law M ∝ 1/(T - T_c) above T_c
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from IPython.display import display

B_FIXED = 0.05
T_HIGH, T_LOW = 5.0, 1.4
N_T = 35
BATCH_SIZE, L = 32, 32
EQUIL_SWEEPS = 50

# Scan from high T toward low T
T_values = np.linspace(T_HIGH, T_LOW, N_T)
sim = IsingSimulator(batch_size=BATCH_SIZE, system_size=L, T=T_HIGH, h=B_FIXED)
M_list, M_err_list = [], []
for T in T_values:
    sim.T = T
    sim.h = B_FIXED
    sim.equilibrate(sweeps=EQUIL_SWEEPS)
    m, m_std = sim.magnetization
    M_list.append(m)
    M_err_list.append(m_std)

M_values = np.array(M_list)
M_err = np.array(M_err_list)

# Curie-Weiss: M = C / (T - T_c) for T > T_c (paramagnetic phase)
def curie_weiss(T, C, Tc):
    return C / (T - Tc)

# Fit only high-T points (well above T_c ≈ 2.27)
mask_fit = T_values > 2.6
T_fit = T_values[mask_fit]
M_fit = M_values[mask_fit]
sigma_fit = np.maximum(M_err[mask_fit], 1e-6)
p0 = (0.5, 2.27)  # initial C, T_c
bounds = ([0, 1.5], [np.inf, 2.5])  # C>0, T_c in [1.5, 2.5]
popt, pcov = curve_fit(curie_weiss, T_fit, M_fit, p0=p0, bounds=bounds, sigma=sigma_fit, absolute_sigma=False)
C_fit, Tc_fit = popt

# Plot
fig, ax = plt.subplots(figsize=(5, 3))
fig.canvas.header_visible = False
ax.set_title(f'M vs T at B={B_FIXED}: Curie-Weiss law M ∝ 1/(T−T_c)')
ax.errorbar(T_values, M_values, yerr=M_err, fmt='o', capsize=2, markersize=5, label='Simulation')
T_smooth = np.linspace(1.01*Tc_fit, T_HIGH, 100)
ax.plot(T_smooth, curie_weiss(T_smooth, *popt), '--', color='C1', lw=2,
        label=f'Curie-Weiss fit: T$_c$={Tc_fit:.2f}')
ax.axvline(Tc_fit, color='gray', ls=':', alpha=0.7, label=f'T$_c$ ≈ {Tc_fit:.2f}')
ax.set_xlabel('T')
ax.set_ylabel('M')
ax.legend(fontsize=8)
ax.set_xlim(T_LOW - 0.1, T_HIGH + 0.1)
ax.set_ylim(-0.05, 1.05)
fig.tight_layout()
display(fig)
plt.close(fig)