import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
from scipy.constants import epsilon_0, mu_0, elementary_charge, Boltzmann

# Constants (using scipy for accuracy)
e = elementary_charge  # 1.60217662e-19 C
kB = Boltzmann  # 1.380649e-23 J/K
eps0 = epsilon_0  # 8.854187817e-12 F/m
mu0 = mu_0  # 1.25663706212e-06 H/m (though document uses 4pi e-7)

# But document uses ε0 = 8.854 × 10^{-12}, e=1.602×10^{-19}
# To match exactly, override if needed
eps0 = 8.854e-12
e = 1.602e-19

# Parameters from document
params = {
    'Parameter': [
        'ne (electron density range)', 'Te (electron temperature range)', 'Ti', 'B (magnetic field)',
        'R (radius)', 'V (potential)', 'md (dust mass range)', 'Qchem (chemical energy)',
        'ηvis (visible efficiency)', 'ρ (air density)', 'U (velocity range)',
        'µeff (effective viscosity range)', "u' (turbulent velocity range)", "ℓ' (length scale range)",
        'ρdust (runaway threshold)', 'χ (ionization potential example)'
    ],
    'Value': [
        '1e16 to 1e18 m^{-3}', '1 to 2 eV', '≈ Te', '50e-6 T',
        '0.1 m', '5 V', '3e-5 to 4e-4 kg', '2e7 J/kg',
        '0.2 to 0.4', '1.2 kg/m^3', '0.1 to 1 m/s',
        '1e-5 to 1e-3 Pa·s', '0.01 to 0.1 m/s', '0.01 to 0.1 m',
        '≥ 0.1 kg/m^3', '15 eV'
    ],
    'Units': [
        'm^{-3}', 'eV', 'eV', 'T',
        'm', 'V', 'kg', 'J/kg',
        '-', 'kg/m^3', 'm/s',
        'Pa·s', 'm/s', 'm',
        'kg/m^3', 'eV'
    ]
}

# Create parameters_table.csv
df_params = pd.DataFrame(params)
df_params.to_csv('parameters_table.csv', index=False)

# Create directories
os.makedirs('calculation_scripts', exist_ok=True)
os.makedirs('output_tables', exist_ok=True)

# Write a sample calculation script to calculation_scripts/main_calculations.py
script_content = """
import numpy as np
import pandas as pd

# Constants
e = 1.602e-19
eps0 = 8.854e-12
mu0 = 4 * np.pi * 1e-7
kB_eV = 1.602e-19  # for Te in eV, kB Te = Te * kB_eV

# Function to compute Debye length
def debye_length(ne, Te_eV):
    return np.sqrt(eps0 * kB_eV * Te_eV / (ne * e**2))

# Function to compute plasma beta
def plasma_beta(ne, Te_eV, Ti_eV, B):
    return 2 * mu0 * ne * kB_eV * (Te_eV + Ti_eV) / B**2

# Function for surface charge
def surface_charge(R, V):
    return 4 * np.pi * eps0 * R * V

# Function for chemical energy
def chemical_energy(md, Qchem):
    return md * Qchem

# Function for radiated power estimate
def prad_eta(md, Qchem, tlife, eta=1):
    return eta * md * Qchem / tlife

# Function for Reynolds number
def reynolds(rho, U, R, mu_eff):
    return rho * U * R / mu_eff

# Function for turbulent diffusion
def dturb(u_prime, l_prime):
    return u_prime * l_prime

# Function for fade time
def tau_fade(R, Dturb):
    return R**2 / Dturb

# Function for Coulomb energy
def e_coulomb(Q, R):
    return Q**2 / (4 * np.pi * eps0 * R)

# Example usage
ne_values = np.logspace(16, 18, num=3)  # 1e16, 1e17, 1e18
Te_values = [1, 2]
B = 50e-6
R = 0.1
V = 5
md_min, md_max = 3e-5, 4e-4
Qchem = 2e7
tlife_typ = 1  # s, example for estimate

results = []
for ne in ne_values:
    for Te in Te_values:
        Ti = Te
        lambda_D = debye_length(ne, Te)
        beta = plasma_beta(ne, Te, Ti, B)
        results.append({'ne': ne, 'Te': Te, 'lambda_D (m)': lambda_D, 'beta': beta})

df_results = pd.DataFrame(results)
df_results.to_csv('output_tables/results_table.csv', index=False)

# Energy balance example
Echem_min = chemical_energy(md_min, Qchem)
Echem_max = chemical_energy(md_max, Qchem)
print(f'Chemical energy: {Echem_min} to {Echem_max} J')

# Prad example (assuming eta=1, tlife=40-1600 s to get 5-15W, but adjust)
# Actually, from 600/120=5W, 8000/533≈15W, but typical 5-15W
Prad_typ = [5, 15]  # W

# Lifetime estimate from tau_fade
Dturb_typ = 1e-3  # m2/s
tau_fade_typ = tau_fade(R, Dturb_typ)
print(f'Typical fade time: {tau_fade_typ} s')

# For spectral synthesis, Saha equation example
def saha_ratio(ni_plus1_over_ni, ne, Te_eV, chi, g_ratio=2):
    return g_ratio * (2 * np.pi * 9.109e-31 * kB_eV * Te_eV / (6.626e-34)**2 )**(3/2) / ne * np.exp(-chi / Te_eV)

# Example, but needs specific chi, etc.
"""

with open('calculation_scripts/main_calculations.py', 'w') as f:
    f.write(script_content)

# Now, run some calculations and produce outputs
# Grid for ne, Te
ne_values = np.array([1e16, 1e17, 1e18])
Te_values = np.array([1, 2])

# Compute lambda_D and beta
lambda_D = np.zeros((len(ne_values), len(Te_values)))
beta = np.zeros((len(ne_values), len(Te_values)))

def debye_length(ne, Te_eV):
    return np.sqrt(eps0 * (Te_eV * e) / (ne * e**2))

def plasma_beta(ne, Te_eV, B=50e-6):
    Ti_eV = Te_eV
    return 2 * mu0 * ne * (Te_eV + Ti_eV) * e / B**2

for i, ne in enumerate(ne_values):
    for j, Te in enumerate(Te_values):
        lambda_D[i, j] = debye_length(ne, Te)
        beta[i, j] = plasma_beta(ne, Te)

# Create DataFrame
data = []
for i, ne in enumerate(ne_values):
    for j, Te in enumerate(Te_values):
        data.append({
            'ne (m^-3)': ne,
            'Te (eV)': Te,
            'lambda_D (m)': lambda_D[i, j],
            'lambda_D (um)': lambda_D[i, j] * 1e6,
            'beta': beta[i, j]
        })

df_output = pd.DataFrame(data)
df_output.to_csv('output_tables/plasma_parameters.csv', index=False)

# Example for Prad and tlife
md_values = [3e-5, 4e-4]
Qchem = 2e7
Echem = [md * Qchem for md in md_values]  # 600, 8000 J
Prad_range = [5, 15]  # W
tlife_estimates = [Echem[0]/Prad_range[1], Echem[1]/Prad_range[0]]  # rough min max

# But to make table
prad_data = [{'Prad (W)': prad} for prad in np.linspace(5,15,5)]
df_prad = pd.DataFrame(prad_data)
df_prad.to_csv('output_tables/radiated_power.csv', index=False)

# Plot example for lambda_D vs ne
fig, ax = plt.subplots()
for Te in Te_values:
    ax.loglog(ne_values, lambda_D[:, Te_values.tolist().index(Te)] * 1e6, label=f'Te={Te} eV')
ax.set_xlabel('ne (m^{-3})')
ax.set_ylabel('lambda_D (um)')
ax.legend()
plt.savefig('output_tables/debye_length_plot.png')

# Similar for beta
fig, ax = plt.subplots()
for Te in Te_values:
    ax.loglog(ne_values, beta[:, Te_values.tolist().index(Te)], label=f'Te={Te} eV')
ax.set_xlabel('ne (m^{-3})')
ax.set_ylabel('beta')
ax.legend()
plt.savefig('output_tables/beta_plot.png')

# For uncertainty scan, assume +/- 10% on ne, Te
# Example for lambda_D at nominal ne=1e17, Te=1
nominal_lambda = debye_length(1e17, 1)
ne_unc = np.linspace(1e17*0.9, 1e17*1.1, 5)
Te_unc = np.linspace(0.9, 1.1, 5)
unc_data = []
for ne_u in ne_unc:
    for Te_u in Te_unc:
        unc_data.append({'ne_unc': ne_u, 'Te_unc': Te_u, 'lambda_D_unc': debye_length(ne_u, Te_u)})

df_unc = pd.DataFrame(unc_data)
df_unc.to_csv('output_tables/uncertainty_scan.csv', index=False)

# Print completion
print("Files generated:")
print("- parameters_table.csv")
print("- calculation_scripts/main_calculations.py")
print("- output_tables/plasma_parameters.csv")
print("- output_tables/radiated_power.csv")
print("- output_tables/debye_length_plot.png")
print("- output_tables/beta_plot.png")
print("- output_tables/uncertainty_scan.csv")