#======================================================================================
# This script is an example showing how to use the LineAnalysisScripting module with : 
# - two molecules (o-H2CO & p-H2CO)
# - two components with the radex model
# - chi2 computation with the Monte-Carlo Markov Chain algorithm
#
# Data files in myDirInput : Orion_Synthetic_H2CO.lis, Orion_Synthetic_Continuum
# Files created in myDirOutput : myName_bestModel.lis, myName.lam, myName.dat, myName_corner_1.png, myName_corner_2.png, myName.txt (statistics)
# (see below for the definitions of myDirInput, myDirOutput and myName)
# 
# To run this script as it is, the only thing you need to specify is the path to your python executable (myPython variable).
# This script might take some time to converge.
#
# Last update: 12-jan-2026
# Tested on: CASSIS 6.5
#
# CASSIS scripts and scripting documentation :
#   - http://cassis.irap.omp.eu/docs/CASSIS_scripts_summary.pdf
#   - http://cassis.irap.omp.eu/docs/script/README.html
#======================================================================================

import os
import time
import subprocess
import ScriptEnvironment
from Range import Range
from Component import Component
from FileReader import FileReader
from LineAnalysisScripting import UserInputs
from eu.omp.irap.cassis.properties import Software
from cassisStats import writeStats
from java.io import File
from Plot import Plot


# INPUTS
                                    
# Define the velocity ranges for the lines
Range.unit = "km/s"  # Possible units are GHz and km/s, MHz, cm-1 and micrometer
v1 = Range(0.0, 12.0)

# Define the step of the walkers values in the MCMC, to be tuned to keep the acceptance rate within [0.2, 0.5]
rpp = 35

# Set some useful variables
sourceName = "Orion_Synthetic"
speciesName = "o-p-H2CO"
myModel = "2C_radex"
myName = sourceName + "_" + speciesName + "_" + myModel

# Change the two lines below for your own myDirInput & myDirOutput folder for your own scripts
myDirInput = Software.getCassisPath()+"/delivery/script/data/"
myDirOutput = Software.getCassisPath()+"/delivery/script/outputs/"
if not os.path.exists(myDirOutput):
    os.makedirs(myDirOutput)

inputFile = myDirInput + "Orion_Synthetic_H2CO.lis"
outputFile = myDirOutput + myName + ".dat"

# This scripts calls an external python script (see part D.) ; set the appropriate path for your python executable below
# (see requirements.txt in CASSIS_PATH/delivery/script/python/triangle_plots/ for the required packages)
myPython = "~/PycharmProjects/cassis-scripts/.venv/bin/python3"
myPythonScript = Software.getCassisPath()+"/delivery/script/python/triangle_plots/plots_MCMC.py"

# CONFIG

userInputs = UserInputs(
    inputFile = inputFile,
    telescope = "hifi",
    tuningRange = [480, 640],
    tuningBand = 60,
    aijMin = 1.0e-4,
    eup = [0.0, 250.0],
    kup = ["*","*"],
    template = "Full VASTEL",
    moltags = [30591,30581],
    tmb2ta = False,
    isoUnique = False,
    outputFile = outputFile,
    plotTitle = myName,
    warning = True,

    # Enter here the lines and the corresponding ranges to be taken into account in the computation
    # The lines should be sorted by increasing frequency.
    selectedLines = {"1-2,4-15": v1},

    # rms of the data around each selected line (in K)
    rmsLines = {"1-2,4-15": 0.025},

    # calibration accuracy of the data around each selected line
    calLines = {"1-2,4-15": 0.15},
    continuum = myDirInput + "Orion_Synthetic_Continuum"
)

# MODEL INPUTS

comp_1 = Component(
    model = "radex",
    # Parameters needed for LTE or RADEX model
    nmol = {'min': 1.0e14, 'max': 5.0e14, 'nstep': 1, 'log_mode': False},
    temp = {'min': 100.0, 'max': 100.0, 'nstep': 1, 'log_mode': False},
    # warning: for radex, use a T range within the T range valid for the collision file
    fwhm = {'min': 5.0, 'max': 5.0, 'nstep': 1, 'log_mode': False},
    size = {'min': 10, 'max': 10, 'nstep': 1, 'log_mode': False},
    vlsr = {'min': 6.0, 'max': 6.0, 'nstep': 1, 'log_mode': False},
    iso = {'min': 1, 'max': 5, 'nstep': 1, 'log_mode': False},  # if moltags are two isotopologues
    interacting = True,
    reducePhysicalParam = {"nmol": rpp, "temp": rpp, "fwhm": rpp, "size": rpp, "vlsr": rpp, "iso": rpp, "n_p_H2": rpp, "n_o_H2": rpp},
    # Parameters below are only needed for RADEX
    collisioner = ["p_H2","o_H2"],
    n_p_H2 = {'min': 2.0e7, 'max': 2.0e7, 'nstep': 1, 'log_mode': False},  # density for the collisioner(s)
    n_o_H2 = {'min': 7.0e7, 'max': 7.0e7, 'nstep': 1, 'log_mode': False},  # density for the collisioner(s)
    collisionFile = ["o-H2CO-o-p-H2-cdms.dat", "p-H2CO-o-p-H2-cdms.dat"],
    geometry = "slab"
)

comp_2 = Component(
    model = "radex",
    # Parameters needed for LTE or RADEX model
    nmol = {'min': 1.0e13, 'max': 5.0e13, 'nstep': 1, 'log_mode': False},
    temp = {'min': 20.0, 'max': 20.0, 'nstep': 1, 'log_mode': False},
    # warning: for radex, use a T range within the T range valid for the collision file
    fwhm = {'min': 1.0, 'max': 1.0, 'nstep': 1, 'log_mode': False},
    size = {'min': 100, 'max': 100, 'nstep': 1, 'log_mode': False},
    vlsr = {'min': 7.0, 'max': 7.0, 'nstep': 1, 'log_mode': False},
    iso = {'min': 3, 'max': 3, 'nstep': 1, 'log_mode': False},  # if moltags are two isotopologues
    interacting = True,
    reducePhysicalParam = {"nmol": rpp, "temp": rpp, "fwhm": rpp, "size": rpp, "vlsr": rpp, "iso": rpp, "n_p_H2": rpp, "n_o_H2": rpp},
    # Parameters below are only needed for RADEX
    collisioner = ["p_H2", "o_H2"],
    n_p_H2 = {'min': 1.0e5, 'max': 1.0e7, 'nstep': 1, 'log_mode': False},  # density for the collisioner(s)
    n_o_H2 = {'min': 3.0e5, 'max': 3.0e7, 'nstep': 1, 'log_mode': False},  # density for the collisioner(s)
    collisionFile = ["o-H2CO-o-p-H2-cdms.dat", "p-H2CO-o-p-H2-cdms.dat"],
    geometry = "slab"
)

# Set the initial parameters by specifying the keys and the values
# Empty dictionary means random initial parameter
params_1 = {"nmol": 1.4e14, "fwhm": 5.0, "vlsr": 6.0, "temp": 100.0, "size": 10, "iso": 3, "n_p_H2": 2.0e7, "n_o_H2": 7.0e7}
params_2 = {"nmol": 2.4e13, "fwhm": 1.0, "vlsr": 7.0, "temp": 20.0, "size": 100, "iso": 3, "n_p_H2": 2.0e6, "n_o_H2": 1.0e7}

# set the walker and burning values
drawNumber = 1000
cutOff = 1
ratioAtCutOff = 1

# Execution time beginning
timeStart = time.time()

# Initialisation of the MCMC components
userInputs.initComponentsForMCMC([comp_1, params_1], [comp_2, params_2])

# Computation of the minimum chi2 using the MCMC algorithm
userInputs.computeChi2MinUsingMCMC(drawNumber, cutOff, ratioAtCutOff)

# ANALYSIS OF THE RESULTS

# A. Plot the best model and save the corresponding spectra and config files
lineModel = userInputs.plotBestModel(moltag = [30591, 30581], overSampling=3, tuningBand=30, telescope="hifi")
lineModel.saveConfig(File(myDirOutput + myName + ".lam"))
bestPhysicalModels  = userInputs.getBestPhysicalModels()
userInputs.saveBestPhysicalModels(myDirOutput + myName + "_bestModel.lis")

# B. Compute and write the statistics of the parameters
# (for more details, print the documentation with 'print writeStats.__doc__')
writeStats(userInputs, sigmaClip=3)

# C. Plot the acceptance rates
# (print documentation for "Plot" with "print Plot.__doc__", documentation for "FileReader" coming soon)
reader = FileReader(userInputs.outputFile)
reader.read_file(lines_to_skip=1)
columns = reader.columns
iterations = columns['#']
rates = columns['rate']
sp = Plot(sourceName)
sp.plot(
    x=iterations, y=rates, xlabel="Walker", ylabel="Acceptance rate", legend="Acceptance rate",
    lineStyle="line", lineWidth=1.5, plotType="line", lineColor="blue"
)

# CASSIS Execution time ending
timeEnd = time.time()
print("CASSIS execution time = %.3f seconds." % (timeEnd - timeStart))

# D. Launch the triangle plot in python
# Set the fraction of rejected walkers here !
fracOfRejWalkers = "0.1"
trianglePlot = [myPython + " " + myPythonScript + " " + myDirOutput + " " + myName + " " + fracOfRejWalkers]
subprocess.Popen(trianglePlot, shell=True)
