# -*- coding: utf-8 -*-

"""
/***************************************************************************
 PipelineEngineer
                                 A QGIS plugin
 A collection of Pipeline Engineering Tools.
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2025-12-06
        copyright            : (C) 2025 by Pipeline Engineer
        email                : info@pipelineengineer.com.au
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

__author__ = 'Pipeline Engineer'
__date__ = '2025-12-06'
__copyright__ = '(C) 2025 by Pipeline Engineer'

# This will get replaced with a git SHA1 when you do a git archive

__revision__ = '$Format:%H$'

from qgis import processing

from qgis.PyQt.QtCore import QCoreApplication
from qgis.core import (
    QgsProcessing,
    QgsFeatureSink,
    QgsProcessingAlgorithm,
    QgsProcessingParameterFeatureSource,
    QgsProcessingParameterFeatureSink,
    QgsProcessingParameterField,
    QgsFeatureRequest,
    QgsProcessingParameterBoolean,
    QgsProcessingParameterNumber,
    QgsProcessingParameterVectorLayer,
    QgsProcessingParameterMultipleLayers,
    QgsProcessingParameterEnum,
    QgsProcessingOutputString,
    QgsProcessingException,
    QgsProcessingOutputMultipleLayers,
    QgsProcessingParameterFileDestination,
    QgsProcessingParameterString,
    QgsProcessingParameterRasterLayer
)

from PyQt5.QtCore import QCoreApplication

from qgis.PyQt.QtGui import QIcon

from .logic.tp_node_correct_directionality import *
from .logic.two_phase_nodal_pipeflow import *

class NodalTwoPhaseFlowAlgorithm(QgsProcessingAlgorithm):

    INPUT = 'INPUT'
    OUTPUT = 'OUTPUT'
    LAYERS = 'LAYERS'
    LIQUID = 'LIQUID'
    GAS = 'GAS'
    CORRELATION = 'CORRELATION'
    SURF_TENS = 'SURF_TENS'
    AMBIENT_TEMP = 'AMBIENT_TEMP'
    FLUID_PRES = 'FLUID_PRES'
    K_H = 'K_H'
    APPLY_HENRYS_LAW = 'APPLY_HENRYS_LAW'
    CORRECT_SALINITY = 'CORRECT_SALINITY'
    C_SALT = 'C_SALT'
    SALINITY = 'SALINITY'
    SALT_MOLAR_MASS = 'SALT_MOLAR_MASS'
    CALC_MODE = 'CALC_MODE'
    MAX_ITER_HYD = 'MAX_ITER_HYD'
    MAX_ITER_THERM = 'MAX_ITER_THERM'
    PRES_TOL = 'PRES_TOL'
    VEL_TOL = 'VEL_TOL'
    TEMP_TOL = 'TEMP_TOL'
    RES_TOL = 'RES_TOL'
    FRIC_MODEL = 'FRIC_MODEL'
    ALPHA = 'ALPHA'
    NONLINEAR_METHOD = 'NONLINEAR_METHOD'
    UPDATE_HYD_ONLY = 'UPDATE_HYD_ONLY'
    CHECK_CONNECTIVITY = 'CHECK_CONNECTIVITY'
    QUIT_ON_INCONSISTENCY = 'QUIT_ON_INCONSISTENCY'
    USE_NUMBA = 'USE_NUMBA'
    LOAD_LAYERS = 'LOAD_LAYERS'

    def initAlgorithm(self, config):

        self.addParameter(
            QgsProcessingParameterMultipleLayers(
                name=self.LAYERS,
                description='Select Network Layers',
                layerType=QgsProcessing.TypeVector
            )
        )
    
        fluids = ["hgas","lgas","hydrogen","methane","water",
                     "biomethane_pure","biomethane_treated","air"]

        gases = ["hgas","lgas","hydrogen","methane","biomethane_pure","biomethane_treated","air"]

        script_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
        user_fluids_path = os.path.join(script_dir, 'user_settings','user_fluids.csv')

        user_fluids_df = pd.read_csv(user_fluids_path)
        
        user_fluids_list = user_fluids_df['name'].tolist()
        
        fluids.extend(user_fluids_list)

        liquids_df = user_fluids_df[user_fluids_df['is_gas']==False]
        gas_df = user_fluids_df[user_fluids_df['is_gas']==True]
        

        user_liquid_list = liquids_df['name'].tolist()
        user_gas_list = gas_df['name'].tolist()
        
        user_fluids_list = user_fluids_df['name'].tolist()

        liquids=["water"]
        
        liquids.extend(user_liquid_list)
        gases.extend(user_gas_list)

        self.addParameter(
            QgsProcessingParameterEnum(
                name=self.LIQUID,
                description='Select Liquid Phase',
                options=liquids,
                defaultValue=0  
            )
        )

        self.addParameter(
            QgsProcessingParameterEnum(
                name=self.GAS,
                description='Select Gas Phase',
                options=gases,
                defaultValue=6  
            )
        )

        self.addParameter(QgsProcessingParameterNumber(
            self.AMBIENT_TEMP, 'Ambient Temperature',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=100000, defaultValue=308.15
        ))

        self.addParameter(
            QgsProcessingParameterEnum(
                name=self.CORRELATION,
                description='Select Pressure Drop Correlation',
                options=['Homogeneous','Lockhart-Martinelli','Beggs & Brill'],
                defaultValue=0
            )
        )

        self.addParameter(QgsProcessingParameterNumber(
            self.SURF_TENS, "Liquid Surface Tension",
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=10000, defaultValue=0.072
        ))
        
        self.addParameter(QgsProcessingParameterBoolean(
            self.APPLY_HENRYS_LAW, 'Apply Henry\'s Law?',
            defaultValue=False
        ))
        
        
        self.addParameter(QgsProcessingParameterNumber(
            self.K_H, "Henry's Constant",
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1000, defaultValue=0.00118
        ))
        
        self.addParameter(QgsProcessingParameterBoolean(
            self.CORRECT_SALINITY, 'Correct for Salinity?',
            defaultValue=False
        ))
        
        self.addParameter(QgsProcessingParameterNumber(
            self.C_SALT, "Setschenow Constant",
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=10000, defaultValue=0.12
        ))
        
        self.addParameter(QgsProcessingParameterNumber(
            self.SALINITY, "Salinity (PPM)",
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1000000000, defaultValue=2000
        ))
        
        self.addParameter(QgsProcessingParameterNumber(
            self.SALT_MOLAR_MASS, "Salt Molar Mass",
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1000000000, defaultValue=58.44
        ))
        
        self.addParameter(QgsProcessingParameterNumber(
            self.FLUID_PRES, 'Average Mixture Pressure (Used for Calculating Phase Properties)',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=100000, defaultValue=1
        ))

        self.addParameter(
            QgsProcessingParameterEnum(
                name=self.CALC_MODE,
                description='Calculation Mode',
                options=['hydraulics'],
                defaultValue=0  
            )
        )

        self.addParameter(QgsProcessingParameterNumber(
            self.MAX_ITER_HYD, 'Maximum Iterations (Hydraulic)',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1e5, defaultValue=10000
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.MAX_ITER_THERM, 'Maximum Iterations (Thermal)',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1e5, defaultValue=10000
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.PRES_TOL, 'Pressure Error Tolerance',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1, defaultValue=1e-4
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.VEL_TOL, 'Velocity Error Tolerance',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1, defaultValue=1e-4
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.TEMP_TOL, 'Temperature Error Tolerance',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1, defaultValue=1e-4
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.RES_TOL, 'Residual Tolerance',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1, defaultValue=1e-3
        ))

        self.addParameter(
            QgsProcessingParameterEnum(
                name=self.FRIC_MODEL,
                description='Friction Model',
                options=['nikuradse', 'colebrook'],
                defaultValue=1
            )
        )

        self.addParameter(QgsProcessingParameterNumber(
            self.ALPHA, 'Alpha',
            type=QgsProcessingParameterNumber.Double,
            minValue=0, maxValue=1e5, defaultValue=1
        ))

        self.addParameter(
            QgsProcessingParameterEnum(
                name=self.NONLINEAR_METHOD,
                description='Nonlinear Method',
                options=['constant', 'automatic'],
                defaultValue=0
            )
        )

        self.addParameter(QgsProcessingParameterBoolean(
            self.UPDATE_HYD_ONLY, 'Update Hydraulic Matrix Only?',
            defaultValue=False
        ))

        self.addParameter(QgsProcessingParameterBoolean(
            self.CHECK_CONNECTIVITY, 'Check Connectivity?',
            defaultValue=True
        ))

        self.addParameter(QgsProcessingParameterBoolean(
            self.QUIT_ON_INCONSISTENCY, 'Quit on Inconsistency?',
            defaultValue=False
        ))

        self.addParameter(QgsProcessingParameterBoolean(
            self.USE_NUMBA, 'Use Numba?',
            defaultValue=True
        ))

        self.addParameter(QgsProcessingParameterBoolean(
            self.LOAD_LAYERS, 'Load Layers?',
            defaultValue=True
        ))

        # Output: list of layer IDs
        self.addParameter(QgsProcessingParameterFileDestination(
                self.OUTPUT,
                'Output GeoPackage',
                fileFilter='GeoPackage (*.gpkg)'
            )
        )

    # -------------------------------------------------------
    # PROCESS
    # -------------------------------------------------------

    def processAlgorithm(self, parameters, context, feedback):

        # MULTIPLE LAYERS
        layers = self.parameterAsLayerList(parameters, self.LAYERS, context)

        if not layers:
            raise QgsProcessingException("No valid input layers selected.")
        # Extract all parameters
        mode_index = self.parameterAsInt(parameters, self.CALC_MODE, context)
        max_iter_hyd = self.parameterAsDouble(parameters, self.MAX_ITER_HYD, context)
        max_iter_therm = self.parameterAsDouble(parameters, self.MAX_ITER_THERM, context)
        tol_p = self.parameterAsDouble(parameters, self.PRES_TOL, context)
        tol_m = self.parameterAsDouble(parameters, self.VEL_TOL, context)
        tol_T = self.parameterAsDouble(parameters, self.TEMP_TOL, context)
        surf_tens = self.parameterAsDouble(parameters, self.SURF_TENS, context)
        tol_res = self.parameterAsDouble(parameters, self.RES_TOL, context)
        ambient_temp = self.parameterAsDouble(parameters, self.AMBIENT_TEMP, context)
        friction_model_index = self.parameterAsInt(parameters, self.FRIC_MODEL, context)
        alpha = self.parameterAsDouble(parameters, self.ALPHA, context)
        nonlinear_method_index = self.parameterAsInt(parameters, self.NONLINEAR_METHOD, context)
        only_update = self.parameterAsBool(parameters, self.UPDATE_HYD_ONLY, context)
        check_conn = self.parameterAsBool(parameters, self.CHECK_CONNECTIVITY, context)
        quit_incons = self.parameterAsBool(parameters, self.QUIT_ON_INCONSISTENCY, context)
        use_numba = self.parameterAsBool(parameters, self.USE_NUMBA, context)
        
        load_layers = self.parameterAsBool(parameters, self.LOAD_LAYERS, context)


        fluids=["hgas","lgas","hydrogen","methane","water","biomethane_pure","biomethane_treated","air"]

        liquids=["water"]

        gases=["hgas","lgas","hydrogen","methane","biomethane_pure","biomethane_treated","air"]

        correlations = ['homogeneous','lockhart_martinelli','beggs_brill']

        fluids = ["hgas","lgas","hydrogen","methane","water",
                     "biomethane_pure","biomethane_treated","air"]

        script_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
        user_fluids_path = os.path.join(script_dir, 'user_settings','user_fluids.csv')

        user_fluids_df = pd.read_csv(user_fluids_path)
        liquids_df = user_fluids_df[user_fluids_df['is_gas']==False]
        gas_df = user_fluids_df[user_fluids_df['is_gas']==True]
        
        user_liquid_list = liquids_df['name'].tolist()
        user_gas_list = gas_df['name'].tolist()
        
        user_fluids_list = user_fluids_df['name'].tolist()
        
        fluids.extend(user_fluids_list)
        
        liquids.extend(user_liquid_list)

        gases.extend(user_gas_list)

        modes=['hydraulics', 'bidirectional', 'sequential']

        apply_henrys_law = self.parameterAsBool(parameters, self.APPLY_HENRYS_LAW, context)
        correct_salinity = self.parameterAsBool(parameters, self.CORRECT_SALINITY, context)

        friction_models =['nikuradse', 'colebrook']

        nonlinear_methods =['constant', 'automatic']

        mode = modes[mode_index]
        friction_model = friction_models[friction_model_index]
        nonlinear_method = nonlinear_methods[nonlinear_method_index]
        
        liquid_phase = liquids[self.parameterAsInt(parameters, self.LIQUID, context)]
        gas_phase = gases[self.parameterAsInt(parameters, self.GAS, context)]

        correlation = correlations[self.parameterAsInt(parameters, self.CORRELATION, context)]
        
        fluid_pres =  self.parameterAsDouble(parameters, self.FLUID_PRES, context)

        k_h = self.parameterAsDouble(parameters, self.K_H, context)
        salt_ppm = self.parameterAsDouble(parameters, self.SALINITY, context)
        c_salt = self.parameterAsDouble(parameters, self.C_SALT, context)
        salt_mol_mass = self.parameterAsDouble(parameters, self.SALT_MOLAR_MASS, context)

        settings = {'max_iter_hyd': max_iter_hyd, 'max_iter_therm': max_iter_therm, 
            'tol_p': tol_p, 'tol_m': tol_m, 'tol_T': tol_T, 'tol_res': tol_res, 
            'ambient_temperature': ambient_temp, 
            'friction_model': friction_model, 'alpha': alpha, 'nonlinear_method': nonlinear_method, 'mode': mode, 
            'only_update_hydraulic_matrix': only_update, 'check_connectivity': check_conn, 
            'quit_on_inconsistency_connectivity': quit_incons, 'use_numba': use_numba}

        gpkg_path = self.parameterAsFileOutput(
                                                parameters,
                                                self.OUTPUT,
                                                context
                                            )

        # Run your solver
        feedback.pushInfo("Running pipeflow solver…")

        result_layers = two_phase_pipeflow_nodal(
                                                layers=layers,
                                                args=settings,
                                                liquid_phase=liquid_phase,gas_phase=gas_phase,
                                                fluid_pres=fluid_pres,fluid_temp=ambient_temp,
                                                k_h=k_h,salt_ppm=salt_ppm,c_salt=c_salt,salt_mol_mass=salt_mol_mass,
                                                surf_tens=surf_tens,apply_henrys_law=apply_henrys_law,correct_salinity=correct_salinity,
                                                correlation=correlation,
                                                feedback=feedback
                                            )
        
        processing.run("native:package", 
                    {'LAYERS':result_layers,
                     'OUTPUT':gpkg_path,
                     'OVERWRITE':False,
                     'SAVE_STYLES':True,
                     'SAVE_METADATA':True,
                     'SELECTED_FEATURES_ONLY':False,
                     'EXPORT_RELATED_LAYERS':False}
                   )

        if load_layers:
            for layer in result_layers:
                QgsProject.instance().addMapLayer(layer)

        return {}
        
    # -------------------------------------------------------

    def name(self):
        return 'nodal_two_phase_pipeflow'

    def displayName(self):
        return self.tr('Two-Phase Flow (Node-to-Node)')

    def group(self):
        return self.tr(self.groupId())

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def shortHelpString(self):
        return self.tr("""
        
    Two-phase fluid solver for calculating pressures from flowrates.
    
    Runs pandapipes pipeflow for primary fluid to determine flow distribution.
    
    Lines with a negative mass flow rate are reversed. Lines with minor flowrates along them (less than 0.000001kg/s) are dropped to avoid errors. An elevation profile at user determined intervals (Chainage) is created for each line.
    
    User can visualise flow distribution and topography of network by ensuring 'Return Network Skeleton' is checked.
    
    Utilising an inital estimate of gas mass per liquid mass, gas void fraction for each user-defined section (chainage) is calculated. The Beggs and Brill correlation is then used to identify the flow regime and concequently calculate the pressure drop.
    
    Starting from the pressure boundary (Grid Layer), the pressure drop along each segment is calculated. Once the end of a line is reached, the "from_junction" of that line is assigned the new calculated pressure drop.
    
    This process is repeated until every node of each line remaining is calculated.
        
        """)

    def createInstance(self):
        return NodalTwoPhaseFlowAlgorithm()

