# python script for generating data for Figure 1 in 
# Parametric resonance, chaos and spatial structure in the Lotka--Volterra model
# Mohamed Swailem and Alastair M. Rucklidge
# Copyright 2025 University of Leeds
# Unless otherwise stated, this dataset is licensed under a Creative Commons 
# Attribution 4.0 International Licence: 
# https://creativecommons.org/licenses/by/4.0/.

# LV ODE mean field stability calculation: 
# run the ODEs for a transient and then see what we have left
# 0 = eqm point (alpha=1 only), 1, 2, 3... = periodic orbit, -1 = chaos

import math
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import cmath

# version=1.00 # 2 May 2025: first working version to do multiple parameter values
# version=1.01   # increased transient integration to 10000, added stroboscopic magenta crosses
# version=1.02   # increased transient integration even more to 100000, printed out stroboscopic points
version=1.03   # longer search for periods (to 99), reduce transient integration to 10000 for alpha=1
               # and save the phase portrait only if the period > 2 (or chaos), or n = 0.7
# edited for publication

pdf_png_extension = '.pdf'   # save as pdf or png
index = 0 # not used for a single calculation

# parameters -- equilibrium point scaled to (astar,1)
# carrying capacity is 1/(kk0 + kk1*cos(omega*time)), where omega=omega0/n
# and omega0 is the imaginary part of the eigenvalue of the (astar,1) equilibrium
# alpha is the homotopy parameter: alpha=0 has time-dependent predation rate

kappa0 = 0.25
kappa1 = 0.24    # instability with alpha=1 is around kappa1=0.225
astar = 1
n = 0.70
alpha = 0.0  # stability calculations need alpha = 1

tol = 1e-8

# derived quantities
lamda0 = (1 - (1 + astar) * kappa0) / astar
# lamda1 = (1 - (1 + astar) * (kappa0 + kappa1*cos(omega*time)) / astar
omega0squared = (1-kappa0)*(1-kappa0)/astar - kappa0*(1-kappa0) - kappa0*kappa0/4
omega0 = math.sqrt(omega0squared)
omega = omega0 / n

# controlling the calculation
n_forcing_periods_transient = 100000
if (abs(alpha - 1.0) < 1.0e-6):
    n_forcing_periods_transient = n_forcing_periods_transient // 10

a_ic = astar * 1.001
b_ic =     1 * 1.001

save_figs_as_pdfs = 1

# construct a filename
file_name_1 = "LV_ODE_"
file_name_2 = "alpha_" + "{:06.4f}".format(alpha) + \
              "_kappa0_" + "{:06.4f}".format(kappa0) + \
              "_kappa1_" + "{:06.4f}".format(kappa1) + \
              "_n_" + "{:06.4f}".format(n) + \
              "_" + "{:04d}".format(index)
file_name_2 = file_name_2.replace(".", "p")
print("File name: ", file_name_2)

def LV(t, ab):
    a, b = ab
    cosomt = math.cos(omega*t)
    kappa = kappa0 + kappa1*cosomt
    lamda1 = (1 - (1 + astar) * kappa) / astar
    lamda = (1-alpha) * lamda0 + alpha * lamda1
    if (lamda<0) or (kappa<0) or (omega0squared<0):
        print("Problem with parameters: lamda, kappa, omega0squared are:", lamda, kappa, omega0squared)
        error

    return [lamda * a*(b - 1), b*(1 - (a + b) * kappa) - lamda * a * b]

def LV_log(t, ab):
    loga, logb = ab
    a = math.exp(loga)
    b = math.exp(logb)
    cosomt = math.cos(omega*t)
    kappa = kappa0 + kappa1*cosomt
    lamda1 = (1 - (1 + astar) * kappa) / astar
    lamda = (1-alpha) * lamda0 + alpha * lamda1
    if (lamda<0) or (kappa<0) or (omega0squared<0):
        print("Problem with parameters: lamda, kappa, omega0squared are:", lamda, kappa, omega0squared)
        error

    return [lamda * (b - 1), (1 - (a + b) * kappa) - lamda * a]

print("astar, kappa0, kappa1, n, alpha:", astar, kappa0, kappa1, n, alpha)
print("lamda0, omega0, omega:", lamda0, omega0, omega)
print("a_ic, b_ic (before transient):", a_ic, b_ic)
print("Filename: ", file_name_2)

# run the mean field for many periods to get over any transients
print("Running transient: ", n_forcing_periods_transient, " periods of the forcing")
sol = solve_ivp(LV, [0,n_forcing_periods_transient * 2 * math.pi / omega], [a_ic, b_ic], method='RK45', rtol=1e-12, atol=1e-12)
mf_stroboscopic_a = sol.y[0]
mf_stroboscopic_b = sol.y[1]
a_ic = mf_stroboscopic_a[-1]
b_ic = mf_stroboscopic_b[-1]
print("a_ic, b_ic (after transient):", a_ic, b_ic)

period = 0

# look for repeats in the stroboscopic map, going forward a further stroboscopic_step steps
stroboscopic_step = 99
if (abs(a_ic - astar) > 100 * tol) and (abs(b_ic - 1) > 100 * tol):
    print("Not an eqm point, hunt for the period")
    period = -1
    t_eval = 2 * math.pi / omega * np.arange(stroboscopic_step)
    sol = solve_ivp(LV, [0,t_eval[-1]], [a_ic, b_ic], t_eval=t_eval, method='RK45', rtol=1e-12, atol=1e-12, max_step=0.1)

    if sol.status != 0:
        print("Status: ", sol.status)
    mf_stroboscopic_time = sol.t
    mf_stroboscopic_a = sol.y[0]
    mf_stroboscopic_b = sol.y[1]
    print(mf_stroboscopic_a)

# hunt back to see if there is periodicity
    for i_hunt_back in range(stroboscopic_step-1):
        print("Checking 0 and ",i_hunt_back + 1," : ", mf_stroboscopic_a[0], mf_stroboscopic_a[i_hunt_back + 1], abs(mf_stroboscopic_a[0] - mf_stroboscopic_a[i_hunt_back + 1]))
        if (abs(mf_stroboscopic_a[0] - mf_stroboscopic_a[i_hunt_back + 1]) < tol):
            period = i_hunt_back + 1
            break

# after the break, we have period = 1, 2, ... if it is periodic
# period = -1: chaotic 
# period =  0: we shouldn't be in this if block

# save a phase portrait and the data, but only if not an equilibrium point
# and save the phase portrait only if the period > 2, or n = 0.7

# generate and save trajectory data for stroboscopic_step steps, or period steps if it is periodic
    t_max = t_eval[-1]
    if period > 0:
        t_max = period * 2 * math.pi / omega
    sol = solve_ivp(LV, [0,t_max], [a_ic, b_ic], method='RK45', rtol=1e-12, atol=1e-12, max_step=0.1)
    time = sol.t
    a = sol.y[0]
    b = sol.y[1]

    n_points = len(time)
    timeab = np.zeros((n_points, 3))
    timeab[:,0] = time
    timeab[:,1] = a
    timeab[:,2] = b
    np.savetxt(file_name_1 + file_name_2 + "_timeab.txt.gz", timeab)

# extract and save stroboscopic data for stroboscopic_step steps, or period steps if it is periodic
    n_points_stroboscopic = len(mf_stroboscopic_time)
    if period > 0:
        n_points_stroboscopic = period
    timeab_stroboscopic = np.zeros((n_points_stroboscopic, 3))
    timeab_stroboscopic[:,0] = mf_stroboscopic_time[0 : n_points_stroboscopic]
    timeab_stroboscopic[:,1] = mf_stroboscopic_a   [0 : n_points_stroboscopic]
    timeab_stroboscopic[:,2] = mf_stroboscopic_b   [0 : n_points_stroboscopic]
    np.savetxt(file_name_1 + file_name_2 + "_timeab_stroboscopic.txt.gz", timeab_stroboscopic)

    print(mf_stroboscopic_a[0:period])

# plot a phase portrait, but only if the period > 2, or n = 0.7 or we are doing only one run
    if (period > 2) or (period < 0) or (abs(n - 0.7) < 1.0e-6) or (one_run_only == 1):
        fig_phaseportrait, fig_phaseportrait_axes = plt.subplots(nrows = 1, ncols = 1, figsize=(10, 10), tight_layout=True)
        fig_phaseportrait_axes.set_xlabel("$a$")
        fig_phaseportrait_axes.set_ylabel("$b$")
        fig_phaseportrait_axes.plot(a, b, 'k', linewidth=2)
        if period > 0:
            fig_phaseportrait_axes.plot(mf_stroboscopic_a[0:period], mf_stroboscopic_b[0:period], 'm+')
        if period == -1:
            fig_phaseportrait_axes.plot(mf_stroboscopic_a, mf_stroboscopic_b, 'm+', linewidth=2)
        fig_phaseportrait.tight_layout()
        if save_figs_as_pdfs == 1:
            fig_phaseportrait.savefig(file_name_1 + file_name_2 + "_phase_portrait" + pdf_png_extension, bbox_inches='tight')
        else:
            plt.show(block = True)
        plt.close(fig_phaseportrait)

# exit the "if not at a fixed point" block
print("Summary: ", file_name_2, " ", \
      "{:06.4f}".format(astar), "{:06.4f}".format(kappa0), "{:06.4f}".format(kappa1), "{:06.4f}".format(n), "{:06.4f}".format(alpha), \
      "{:06.4f}".format(lamda0), "{:06.4f}".format(omega0), "{:06.4f}".format(omega), \
      "{:06.4f}".format(a_ic), "{:06.4f}".format(b_ic), \
      "{:5.2f}".format(version), \
      "{:6d}".format(period))

