main.py

"""
main.py — Main script for rocket engine nozzle simulation (no regenerative cooling).

Calculates exhaust gas flow through the nozzle (N, P, T) including friction.
Prepares a grid of points for export (e.g. to Ansys).

Stages:
  1. Initial ODE integration (isentropic flow)
  2. Convergence loop: full ODE with friction (dQdx=0) until Y converges

Required libraries: numpy, scipy, matplotlib
Usage:
  python main.py                        — interactive mode (select parameter file)
  python main.py params/default.json   — load a specific file
  python main.py --default             — run with built-in default parameters
"""

import sys
import time
czas_start = time.time()
import numpy as np
from scipy.integrate import solve_ivp

from geometry import build_nozzle_geometry
from ode_functions import my_nozzle_ode
from parameters import compute_gas_parameters
from convergence_loop import run_convergence_loop
from param_loader import load_and_select_params, load_params
from results_exporter import export_results
from gas_properties import build_gas_property_arrays, log_property_nodes


def main(param_file=None):
    # =====================================================================
    # 0. WCZYTANIE PARAMETROW
    # =====================================================================
    loaded = None
    selected_filepath = None

    if param_file == '--default':
        loaded = None
        selected_filepath = None
    elif param_file is not None:
        print(f"Loading parameters from: {param_file}")
        loaded, _ = load_params(param_file)
        selected_filepath = param_file
        print(f"  Loaded {len(loaded)} parameters.")
    else:
        loaded, selected_filepath = load_and_select_params()

    if selected_filepath is not None:
        import os as _os
        params_name = _os.path.splitext(_os.path.basename(selected_filepath))[0]
    else:
        params_name = 'default'

    def p(key, default):
        if loaded and key in loaded:
            return loaded[key]
        return default

    # =====================================================================
    # 1. INICJALIZACJA GEOMETRII DYSZY
    # =====================================================================
    print("Building nozzle geometry...")
    R_throat = p('R_throat', 0.01878)
    E_r = p('E_r', 5)
    n_grid = int(p('n_grid', 100))

    xspan, R_grid, A_grid, dA_grid_dx, A_interp, dA_interp = build_nozzle_geometry(
        R_param=R_throat, E_r=E_r, n_grid=n_grid
    )

    n = len(xspan)
    dx = np.mean(np.diff(xspan))

    # Warunek poczatkowy ODE: [N, P, T]
    Y0 = np.array([p('N0', 0.0153), p('P0', 6000000.0), p('T0', 2941.58)])

    # =====================================================================
    # 2. SLOWNIK PARAMETROW
    # =====================================================================
    params = {}
    params['A'] = A_grid.copy()
    params['R'] = R_grid.copy()

    idx_throat = int(np.argmin(params['R']))

    # --- Stale fizyczne ---
    params['eta'] = p('eta', 0.000086742)
    params['epsilon'] = p('epsilon', 0.00005)
    params['D'] = params['R'] * 2
    params['At'] = np.min(params['A'])
    params['Dt'] = 2 * np.min(params['R'])

    # --- Metadane (eksport CSV) ---
    params['c_star'] = p('c_star', 2416.8)
    params['mdot_gas'] = p('mdot_gas', 1.06)

    # =====================================================================
    # INTERPOLACJA WLASCIWOSCI GAZU (PCHIP: komora -> gardziel -> wylot)
    # =====================================================================
    print("Building gas property profiles (PCHIP)...")
    gas_props = build_gas_property_arrays(p, xspan, idx_throat)
    log_property_nodes(gas_props)

    params['gamma_arr'] = gas_props['gamma_arr']
    params['Cpcg_arr'] = gas_props['Cpcg_arr']
    params['Prcg_arr'] = gas_props['Prcg_arr']
    params['molar_mass_arr'] = gas_props['combustion_molar_mass_arr']
    params['Rs_arr'] = gas_props['Rs_arr']

    params['gamma_interp'] = gas_props['gamma_interp']
    params['Cpcg_interp'] = gas_props['Cpcg_interp']
    params['Prcg_interp'] = gas_props['Prcg_interp']
    params['Rs_interp'] = gas_props['Rs_interp']

    params['gamma'] = gas_props['gamma_chamber']

    # Funkcje interpolujace A(x) i dA/dx(x) dla solvera ODE
    params['A_func'] = A_interp
    params['dA_func'] = dA_interp

    # Pre-alokacja wektorow wynikowych
    params['T_aw'] = np.full(n, np.nan)
    params['M'] = np.full(n, np.nan)

    # =====================================================================
    # ETAP 1: WSTEPNA INTEGRACJA ODE (przeplyw izoentropowy)
    # =====================================================================
    print("Stage 1: ODE integration (isentropic flow)...")

    sol1 = solve_ivp(
        fun=lambda x, Y: my_nozzle_ode(x, Y, params),
        t_span=(xspan[0], xspan[-1]),
        y0=Y0,
        method='RK45',
        t_eval=xspan,
        rtol=1e-3,
        atol=1e-6
    )

    YSol = sol1.y.T
    print(f"  ODE1 complete. Points: {YSol.shape[0]}")

    # Oblicz T_aw i M z rozwiazania izoentropowego
    for i in range(n):
        T_aw_i, M_i = compute_gas_parameters(YSol, params, i)
        params['T_aw'][i] = T_aw_i
        params['M'][i] = M_i

    print("  Stage 1 complete.")

    # =====================================================================
    # PETLA KONWERGENCJI: ETAP 2 powtarzany do zbieznosci Y
    # =====================================================================
    solver_max_iter = int(p('max_iterations', 50))
    solver_tol = p('tol', 1e-6)
    solver_relax = p('relax', 0.5)
    solver_mode = str(p('solver_mode', 'convergence'))   # 'convergence' or 'fixed'

    YSol_final, convergence_history, residual_histories = run_convergence_loop(
        xspan=xspan,
        Y0=Y0,
        params=params,
        YSol_init=YSol,
        dx=dx,
        max_iterations=solver_max_iter,
        tol=solver_tol,
        relax=solver_relax,
        mode=solver_mode,
    )

    # Przelicz T_aw i M z koncowego rozwiazania
    for i in range(n):
        T_aw_i, M_i = compute_gas_parameters(YSol_final, params, i)
        params['T_aw'][i] = T_aw_i
        params['M'][i] = M_i

    czas_koniec = time.time()
    czas_wykonania = czas_koniec - czas_start
    print(f"Execution time: {czas_wykonania:.2f} seconds")

    # =====================================================================
    # EKSPORT WYNIKOW DO PLIKU CSV
    # =====================================================================
    converged = (len(convergence_history) == 0
                 or convergence_history[-1] < solver_tol)
    export_results(
        xspan=xspan,
        YSol_final=YSol_final,
        params=params,
        convergence_history=convergence_history,
        params_name=params_name,
        converged=converged,
    )

    print("Simulation completed successfully!")


if __name__ == '__main__':
    if len(sys.argv) > 1:
        main(param_file=sys.argv[1])
    else:
        main()