#!/bin/bash
set -e

# Steady-state open Dicke simulations for four dissipation scenarios.
# Outputs Wigner function arrays to CSV files named 1.csv-4.csv.

python3 <<'PYTHON_SCRIPT'
import os
from pathlib import Path
import numpy as np

from qutip import *
from qutip.piqs import *

# Choose an output directory that works both locally and in the container.
output_dir_candidates = [
    Path(os.getenv("OUTPUT_DIR", "")),
    Path.cwd(),
    Path("/root"),
]
output_dir = None
for candidate in output_dir_candidates:
    if candidate and candidate.exists():
        output_dir = candidate
        break
if output_dir is None:
    output_dir = Path.cwd()

print(f"Saving CSV files to {output_dir}")

# TLS parameters
N = 4
ntls = N
nds = num_dicke_states(ntls)
[jx, jy, jz] = jspin(N)
jp = jspin(N, "+")
jm = jp.dag()
w0 = 1
gE = 0.1
gD = 0.01
gP = 0.1
gCP = 0.1
gCE = 0.1
gCD = 0.1
h = w0 * jz

# photonic parameters
nphot = 16
wc = 1
kappa = 1
ratio_g = 2
g = ratio_g / np.sqrt(N)
a = destroy(nphot)

# TLS liouvillian
system = Dicke(N=N)
system.hamiltonian = h
system.emission = 0
system.dephasing = gD
system.pumping = gP
system.collective_pumping = 0
system.collective_emission = 0
system.collective_dephasing = 0
liouv = system.liouvillian()

# TLS liouvillian 2
system2 = Dicke(N=N)
system2.hamiltonian = h
system2.emission = gE
system2.dephasing = gD
system2.pumping = 0
system2.collective_pumping = 0
system2.collective_emission = 0
system2.collective_dephasing = 0
liouv2 = system2.liouvillian()

# TLS liouvillian 3
system3 = Dicke(N=N)
system3.hamiltonian = h
system3.emission = gE
system3.dephasing = gD
system3.pumping = 0  # gP
system3.collective_pumping = gCP
system3.collective_emission = 0
system3.collective_dephasing = 0
liouv3 = system3.liouvillian()

# TLS liouvillian 4
system4 = Dicke(N=N)
system4.hamiltonian = h
system4.emission = gE
system4.dephasing = gD
system4.pumping = 0
system4.collective_pumping = 0
system4.collective_emission = gCE
system4.collective_dephasing = 0
liouv4 = system4.liouvillian()

# photonic liouvillian
h_phot = wc * a.dag() * a
c_ops_phot = [np.sqrt(kappa) * a]
liouv_phot = liouvillian(h_phot, c_ops_phot)

# identity operators
id_tls = to_super(qeye(nds))
id_phot = to_super(qeye(nphot))

# light-matter superoperator
h_int = g * tensor(a + a.dag(), jx)
liouv_int = -1j * spre(h_int) + 1j * spost(h_int)

# total liouvillians
liouv_sum = super_tensor(liouv_phot, id_tls) + super_tensor(id_phot, liouv)
liouv_tot = liouv_sum + liouv_int

liouv_sum2 = super_tensor(liouv_phot, id_tls) + super_tensor(id_phot, liouv2)
liouv_tot2 = liouv_sum2 + liouv_int

liouv_sum3 = super_tensor(liouv_phot, id_tls) + super_tensor(id_phot, liouv3)
liouv_tot3 = liouv_sum3 + liouv_int

liouv_sum4 = super_tensor(liouv_phot, id_tls) + super_tensor(id_phot, liouv4)
liouv_tot4 = liouv_sum4 + liouv_int

# total operators
jz_tot = tensor(qeye(nphot), jz)
jp_tot = tensor(qeye(nphot), jp)
jm_tot = tensor(qeye(nphot), jm)
jpjm_tot = tensor(qeye(nphot), jp * jm)
nphot_tot = tensor(a.dag() * a, qeye(nds))
adag_tot = tensor(a.dag(), qeye(nds))
a_tot = tensor(a, qeye(nds))

# calculate steady states
rho_ss4 = steadystate(liouv_tot4, method="direct")
nphot_ss4 = expect(nphot_tot, rho_ss4)
psi4 = rho_ss4.ptrace(0)
print("Ensemble 4 is ok")

rho_ss = steadystate(liouv_tot, method="direct")
nphot_ss = expect(nphot_tot, rho_ss)
psi = rho_ss.ptrace(0)

rho_ss2 = steadystate(liouv_tot2, method="direct")
nphot_ss2 = expect(nphot_tot, rho_ss2)
psi2 = rho_ss2.ptrace(0)

rho_ss3 = steadystate(liouv_tot3, method="direct")
nphot_ss3 = expect(nphot_tot, rho_ss3)
psi3 = rho_ss3.ptrace(0)

# calculate Wigner function for photonic states
nx = 1000
xvec = np.linspace(-6, 6, nx)
W = wigner(psi, xvec, xvec)
print("1 ok")
W2 = wigner(psi2, xvec, xvec)
print("2 ok")
W3 = wigner(psi3, xvec, xvec)
print("3 ok")
W4 = wigner(psi4, xvec, xvec)
print("4 ok")

# Save each simulation result to its own CSV file
for idx, W_sim in enumerate([W, W2, W3, W4], start=1):
    out_path = output_dir / f"{idx}.csv"
    np.savetxt(out_path, W_sim, delimiter=",")
    print(f"Saved {out_path}")

PYTHON_SCRIPT