import numpy as np
import matplotlib.pyplot as plt
import time

class ReactionDiffusionModel:
    def __init__(self, Nx=128, Ny=128, P=0.06):
        self.Nx = Nx
        self.Ny = Ny
        self.dx = 1.0
        self.dy = 1.0
        
        self.epsilon = 4.0
        self.beta = 0.9
        self.rho = 0.25
        self.T = 1.0
        self.D_par = 0.1
        self.P = P
        
        self.dt = 0.01
        self.t_max = 2000
        self.steps = int(self.t_max / self.dt)
        
        np.random.seed(42)
        self.x = 0.5 + 0.1 * (np.random.rand(self.Nx, self.Ny) - 0.5)
        
        self.time_points = []
        self.mean_history = []
        self.variance_history = []
        
        self.snapshots = {}
        
        if abs(self.P - 0.12) < 1e-5:
            self.snapshot_times = [20, 40, 60, 80, 100, 500, 1000, 2000]
        else:
            self.snapshot_times = [20, 100, 120, 160, 200, 500, 1000, 2000]

    def laplacian(self, field):
        top = np.roll(field, 1, axis=0)
        bottom = np.roll(field, -1, axis=0)
        left = np.roll(field, 1, axis=1)
        right = np.roll(field, -1, axis=1)
        return (top + bottom + left + right - 4 * field) / (self.dx ** 2)

    def nonlinear_flux_divergence(self, M, field):
        M_ip1 = np.roll(M, -1, axis=0)
        M_im1 = np.roll(M, 1, axis=0)
        M_jp1 = np.roll(M, -1, axis=1)
        M_jm1 = np.roll(M, 1, axis=1)
        
        f_ip1 = np.roll(field, -1, axis=0)
        f_im1 = np.roll(field, 1, axis=0)
        f_jp1 = np.roll(field, -1, axis=1)
        f_jm1 = np.roll(field, 1, axis=1)
        
        term_x = ( (M_ip1 + M)/2.0 * (f_ip1 - field) - 
                   (M + M_im1)/2.0 * (field - f_im1) ) / (self.dx**2)
                   
        term_y = ( (M_jp1 + M)/2.0 * (f_jp1 - field) - 
                   (M + M_jm1)/2.0 * (field - f_jm1) ) / (self.dy**2)
                   
        return term_x + term_y

    def reaction_term(self, x):
        term1 = self.P * (1 - x) * (1 - self.beta * x)
        term2 = x * (1 - self.beta * x) * np.exp(- (2 * self.epsilon / self.T) * x)
        term3 = self.D_par * x * (1 - self.beta)
        return term1 - term2 - term3

    def run(self):
        print(f"Початок симуляції для P={self.P}...")
        print(f"  Цільові знімки на t: {self.snapshot_times}")
        start_time = time.time()
        
        for step in range(self.steps + 1):
            t = step * self.dt
            
            if step % 100 == 0:
                mean_val = np.mean(self.x)
                sq_mean = np.mean(self.x**2)
                variance = sq_mean - mean_val**2
                
                self.time_points.append(t)
                self.mean_history.append(mean_val)
                self.variance_history.append(variance)
            
            for st in self.snapshot_times:
                if abs(t - st) < self.dt / 1.5: 
                    if st not in self.snapshots:
                        self.snapshots[st] = self.x.copy()

            M = self.x * (1 - self.x)
            lap_x = self.laplacian(self.x)
            
            factor = 2 * self.epsilon / self.T
            term_nabla_M_nabla_x = self.nonlinear_flux_divergence(M, self.x)
            term_nabla_M_nabla_lapX = self.nonlinear_flux_divergence(M, lap_x)
            
            Diffusion_total = lap_x - factor * term_nabla_M_nabla_x - factor * (self.rho**2) * term_nabla_M_nabla_lapX
            Reaction = self.reaction_term(self.x)
            
            self.x += self.dt * (Reaction + Diffusion_total)
            self.x = np.clip(self.x, 0, 1)

            if step % 50000 == 0 and step > 0:
                print(f"  Прогрес: {t:.0f}/{self.t_max} a.u.")

        print(f"Симуляцію завершено за {time.time() - start_time:.2f} с.")


def main():
    sim1 = ReactionDiffusionModel(P=0.06)
    sim1.run()
    
    sim2 = ReactionDiffusionModel(P=0.12)
    sim2.run()
    
    print("Виведення графіків...")
    
    plt.figure(figsize=(8, 6))
    plt.plot(sim1.time_points, sim1.mean_history, 'k-', label='P = 0.06')
    plt.plot(sim2.time_points, sim2.mean_history, 'r-', label='P = 0.12')
    plt.xlabel('time [a.u.]')
    plt.ylabel('')
    plt.title('Середня концентрація')
    plt.legend()
    plt.grid(True, linestyle=':', alpha=0.6)
    
    plt.figure(figsize=(8, 6))
    plt.plot(sim1.time_points, sim1.variance_history, 'k-', label='P = 0.06')
    plt.plot(sim2.time_points, sim2.variance_history, 'r-', label='P = 0.12')
    plt.xlabel('time [a.u.]')
    plt.ylabel('<($\delta$x)$^2$>')
    plt.title('Дисперсія')
    plt.legend()
    plt.grid(True, linestyle=':', alpha=0.6)
    
    def show_8_snapshots(sim_obj, title):
        times = sorted(sim_obj.snapshots.keys())
        fig, axes = plt.subplots(2, 4, figsize=(16, 8))
        fig.suptitle(title, fontsize=16)
        
        ax_flat = axes.flatten()
        last_im = None
        
        for i, ax in enumerate(ax_flat):
            if i < len(times):
                t = times[i]
                last_im = ax.imshow(sim_obj.snapshots[t], cmap='hot', origin='lower', interpolation='bilinear', vmin=0, vmax=1)
                ax.set_title(f"t = {t}")
                ax.axis('off')
            else:
                ax.axis('off')
        
        if last_im:
            cbar = fig.colorbar(last_im, ax=axes, orientation='vertical', fraction=0.025, pad=0.04)
            cbar.set_label('x(r)', fontsize=12)
            
    show_8_snapshots(sim2, f"Еволюція для P=0.12")
    
    show_8_snapshots(sim1, f"Еволюція для P=0.06")
    
    plt.show()

if __name__ == "__main__":
    main()