#!/usr/bin/env python3
# pyrotoolbox, a collection of tools to work with PyroScience GmbH data.
# Copyright (C) 2025, Christoph Staudinger
#
# 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 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
""" Module containing functions for oxygen unit conversions.
"""

import numpy as np
from typing import Iterable
import pandas as pd


def vapour_pressure_water(temperature):
    """ Calculates the vapour pressure of water at given Temperature (°C), return in hPa

    #This is Equation (6) from Murray, F.W. 1967. On the computation of saturation vapour pressure.
    J. Applied Meteorology 6: 203-204

    :param temperature: Temperature in °C
    :return: vapour pressure in hPa
    """
    return 6.1078 * np.exp(17.2694 * temperature / (temperature + 237.3))


def calc_oxygen_solubility(temperature, salinity):
    """ Calculate the oxygen solubility in µM for water with given temperature and salinity.

    The formula and constants from "From Garcia 1992: Oxygen solubility in seawater: better fitting equations" are used.

    :param temperature: water temperature (°C)
    :param salinity: salinity of water (g/L)
    :return: oxygen solubility in µM
    """
    # From Garcia 1992: Oxygen solubility in seawater: better fitting equations
    A0 = 2.00856
    A1 = 3.224
    A2 = 3.99063
    A3 = 4.80299
    A4 = 0.978188
    A5 = 1.71069
    B0 = -0.00624097
    B1 = -0.00693498
    B2 = -0.00690358
    B3 = -0.00429155
    C0 = -3.1168e-07

    TS = np.log((298.15 - temperature) / (273.15 + temperature))

    return np.exp(A0 + A1 * TS + A2 * TS ** 2 + A3 * TS ** 3 + A4 * TS ** 4 + A5 * TS ** 5 + salinity *
                  (B0 + B1 * TS + B2 * TS ** 2 + B3 * TS ** 3) + C0 * salinity ** 2) / 0.02241


def calc_pressure_and_water_corrected_pO2(pressure, T, water_sat, percentO2=20.95):
    """ calculate the partial pressure of oxygen in air with a given water saturation and total pressure

    :param pressure: ambient pressure in hPa
    :param T: Temperature in °C
    :param water_sat: water saturation of the air in % (0-100)
    :param percentO2: percentage of oxygen in air in % (0-100). Default 20.95
    :return: partial pressure of oxygen in hPa
    """
    # calc pressure of water vapour
    pH2O = vapour_pressure_water(T) * water_sat / 100
    # calculate remaining partial pressure of dry air
    pAir = pressure - pH2O
    return pAir * percentO2 / 100


def i_only_think_in_hpa(df: pd.DataFrame, m: dict) -> pd.Series:
    """ Convert any oxygen data to hPa.

    Useful for people like me who prefer to look at oxygen data in hPa.
    This function takes the dataframe and metadata dict of the parser functions and returns the oxygen data in hPa.
    If you want to change the conditions of the conversions (e.g. assuming the experiment setup was wrong) use
    the convert_to_hPa function.

    :param df: DataFrame as generated by the parser functions.
    :param m: metadata-dict as generated by the parser functions.
    :return: Series with oxygen data in hPa.
    """
    if isinstance(m['settings']['pressure'], float):  # TODO what happens here with aquaphox logfiles? Can i set pressure dependent units?  how are they calculated?
        pressure = m['settings']['pressure']
    elif m['settings']['pressure'] == 'internal sensor':
        pressure = df['pressure']
    else:
        raise ValueError(f'Failed to parse pressure setting "{m["settings"]["pressure"]}"')
    if isinstance(m['settings']['temperature'], float):
        temperature = m['settings']['temperature']
    elif m['settings']['temperature'].startswith('Optical'):
        temperature = df['optical_temperature']
    elif m['settings']['temperature'] == 'external sensor':
        temperature = df['sample_temperature']
    elif m['settings']['temperature'] == 'internal sensor':
        temperature = df['case_temperature']
    else:
        raise ValueError(f'Failed to parse temperature setting "{m["settings"]["temperature"]}"')
    salinity = m['settings']['salinity']

    for unit in ('oxygen_hPa', 'oxygen_torr', 'oxygen_%O2', 'oxygen_%airsat', 'oxygen_µM', 'oxygen_µg/L',
                 'oxygen_mg/L', 'oxygen_mL/L'):
        if unit in df:
            break
    else:
        raise ValueError('No known oxygen unit found in data!')

    return convert_to_hPa(df[unit], unit, temperature, pressure, salinity)


def i_have_a_fireplate_and_still_only_think_in_hPa(df: pd.DataFrame, m: dict) -> pd.DataFrame:
    """ A copy of the popular "i_only_think_in_hpa" function for fireplate data.

    This was separated due to different format of fireplate data, but works very similar.

    :param df: DataFrame as generated by the read_fireplate_workbench function.
    :param m: metadata-dict as generated by the read_fireplate_workbench function.
    :return: DataFrame with oxygen data in hPa.
    """
    if isinstance(m['settings']['pressure'], float):
        pressure = m['settings']['pressure']
    elif m['settings']['pressure'] == 'internal sensor':
        pressure = df['pressure']
    else:
        raise ValueError(f'Failed to parse pressure setting "{m["settings"]["pressure"]}"')
    if isinstance(m['settings']['temperature'], float):
        temperature = m['settings']['temperature']
    elif m['settings']['temperature'].startswith('Optical'):
        temperature = df['optical_temperature']
    elif m['settings']['temperature'] == 'external sensor':
        temperature = df['sample_temperature']
    elif m['settings']['temperature'] == 'internal sensor':
        temperature = df['case_temperature']
    else:
        raise ValueError(f'Failed to parse temperature setting "{m["settings"]["temperature"]}"')
    salinity = m['settings']['salinity']

    d = pd.DataFrame()
    for c in df.filter(regex='oxygen'):
        for unit in ('oxygen_hPa', 'oxygen_torr', 'oxygen_%O2', 'oxygen_%airsat', 'oxygen_µM', 'oxygen_µg/L',
                     'oxygen_mg/L', 'oxygen_mL/L'):
            if unit in c:
                break
        else:
            raise ValueError(f'No known oxygen unit found in column "{c}"!')

        d[c[:4] + 'oxygen_hPa'] = convert_to_hPa(df[c], unit, temperature, pressure, salinity)
    return d


def convert_to_hPa(data: pd.Series, unit: str, temperature=None, pressure=None, salinity=None) -> pd.Series:
    """Convert any oxygen unit to hPa.

    :param data: any oxygen data produced by a pyroscience device.
    :param unit: unit of the data. Has to be one of ('oxygen_hPa', 'oxygen_torr', 'oxygen_%O2', 'oxygen_%airsat', 'oxygen_µM', 'oxygen_µg/L', 'oxygen_mg/L', 'oxygen_mL/L').
    :param temperature: sample temperature. Used for all units except 'oxygen_hPa', 'oxygen_torr' and 'oxygen_%O2'.
    :param pressure: gas-pressure of or above the sample. Used for all units except 'oxygen_hPa' and 'oxygen_torr'
    :param salinity: salinity of the sample in g/L. Used for all dissolved oxygen units.
    """
    if unit not in ('oxygen_hPa', 'oxygen_torr', 'oxygen_%O2', 'oxygen_%airsat', 'oxygen_µM', 'oxygen_µg/L',
                    'oxygen_mg/L', 'oxygen_mL/L'):
        raise ValueError('No known oxygen unit found in data!')

    if unit == 'oxygen_hPa':
        return data
    elif unit == 'oxygen_torr':
        return data * 1.33322
    if pressure is None and unit in ('oxygen_%O2', 'oxygen_%airsat'):
        raise ValueError('pressure data is missing for this unit conversion!')
    if unit == 'oxygen_%O2':
        return data / 100 * pressure

    if temperature is None:
        raise ValueError('temperature data is missing for this unit conversion')
    ph2o = vapour_pressure_water(temperature)
    if unit == 'oxygen_%airsat':
        p_100percentairsat = (pressure - ph2o) * 0.2095
        return data / 100 * p_100percentairsat

    if salinity is None:
        raise ValueError('salinity data is missing for this unit conversion')
    ox_solubility = calc_oxygen_solubility(temperature, salinity)
    ox_solubility /= (1013 - ph2o) * 0.2095  # correct solubility for water vapor partial pressure

    if unit == 'oxygen_µM':
        return data / ox_solubility
    elif unit == 'oxygen_µg/L':
        return data / 1000 / 31.999 * 1000 / ox_solubility
    elif unit == 'oxygen_mg/L':
        return data / 31.999 * 1000 / ox_solubility
    elif unit == 'oxygen_mL/L':
        return data / 0.02241 / ox_solubility
    raise ValueError('No known oxygen unit found in data!')


def hPa_to_torr(data):
    """Convert hPa to torr.

    :param data: hPa data to convert.
    """
    return data / 1.33322


def hPa_to_percentO2(data, pressure):
    """Convert hPa to %O2.

    :param data: hPa data to convert.
    :param pressure: gas pressure of the sample in hPa.
    """
    return data / pressure * 100


def hPa_to_percent_airsat(data, pressure, temperature):
    """ Convert hPa to %airsat. This is a unit for water samples and assumes 100% relative humidity.

    :param data: hPa data to convert.
    :param pressure: gas pressure above the sample in hPa.
    :param temperature: temperature of the sample in °C.
    """
    ph2o = vapour_pressure_water(temperature)
    p_100percentairsat = (pressure - ph2o) * 0.2095
    return data / p_100percentairsat * 100


def hPa_to_uM(data, temperature, salinity):
    """ Convert hPa to µM. This is a unit for water sample.

    :param data: hPa data to convert.
    :param temperature: temperature of the sample in °C.
    :param salinity: salinity of the sample in g/L.
    """
    ph2o = vapour_pressure_water(temperature)
    ox_solubility = calc_oxygen_solubility(temperature, salinity)
    return data * ox_solubility / (1013 - ph2o) / 0.2095


def hPa_to_mgL(data, temperature, salinity):
    """ Convert hPa to mg/L. This is a unit for water sample.

    :param data: hPa data to convert.
    :param temperature: temperature of the sample in °C.
    :param salinity: salinity of the sample in g/L.
    """
    return hPa_to_uM(data, temperature, salinity) * 31.999 / 1000


def calc_tau(dphi: float, f: float) -> float:
    """ calculate lifetime from phase angle and modulation frequency

    :param dphi: phase angle in °
    :param f: modulation frequency in Hz
    """
    return np.tan(np.radians(dphi))/2/np.pi/f

def _calc_tau0_20_and_ksv_20(dphi0, dphi100, temp0, temp100, pressure, humidity, percentO2, f, m, kt, tt, mt, ft=0,
                             freq=4000) -> (float, float):
    """ calculate tau0 and ksv for 20°C from a calibration dictionary.

    tau0 is returned in µs and ksv in 1/hPa.

    :param dphi0: phase angle at zero calibration point in °
    :param dphi100: phase angle at the upper calibration point in !
    :param temp0: temperature at zero calibration point in °C
    :param temp100: temperature at the upper calibration point in °C
    :param pressure: pressure at the upper calibration point in mbar/hPa
    :param humidity: relative humidity at the upper calibraiton point in %
    :param percentO2: oxygen content in the dry calibration gas (upper calibration point) in %
    :param f: sensor constant
    :param m: sensor constant
    :param kt: sensor constant in 1/K
    :param tt: sensor constant in 1/K
    :param mt: sensor constant in 1/K
    :param ft: sensor constant in 1/K (hardly used)
    :param freq: modulation frequency in Hz
    """
    tau0 = calc_tau(dphi0, freq) * 1e6  # convert to µs
    tau100 = calc_tau(dphi100, freq) * 1e6  # convert to µs
    pO2 = calc_pressure_and_water_corrected_pO2(pressure, temp100, humidity, percentO2)

    tau0_20 = _calc_tau0_20(tau0, temp0, tt)
    ksv_20 = _calc_ksv_20(tau0_20, tau100, pO2, temp100, f, m, kt, tt, mt, ft)
    return tau0_20, ksv_20


def calculate_pO2(dphi, temperature, dphi0, dphi100, temp0, temp100, pressure, humidity, percentO2, f, m, kt, tt, mt,
                  ft=0, freq=4000, **kwargs):
    """ Calculate pO2 in hPa for a given phase-angle and temperature.

    :param dphi: phase angle in °
    :param temperature: in °C
    :param dphi0: phase angle at zero calibration point in °
    :param dphi100: phase angle at the upper calibration point in !
    :param temp0: temperature at zero calibration point in °C
    :param temp100: temperature at the upper calibration point in °C
    :param pressure: pressure at the upper calibration point in mbar/hPa
    :param humidity: relative humidity at the upper calibration point in %
    :param percentO2: oxygen content in the dry calibration gas (upper calibration point) in %
    :param f: sensor constant
    :param m: sensor constant
    :param kt: sensor constant (in 1/K)
    :param tt: sensor constant (in 1/K)
    :param mt: sensor constant (in 1/K)
    :param ft: sensor constant (in 1/K)
    :param freq: modulation frequency in Hz
    :param kwargs: to accept additional unused parameters. (allow passing \*\*calibration without an error)
    :return: pO2 in hPa
    """
    tau0_20, ksv_20 = _calc_tau0_20_and_ksv_20(dphi0, dphi100, temp0, temp100, pressure, humidity, percentO2, f, m, kt,
                                               tt, mt, ft, freq)
    tau = calc_tau(dphi, freq) * 1e6

    return _calc_pO2(tau, temperature, tau0_20, ksv_20, f, m, kt, tt, mt, ft)


def calculate_pO2_from_calibration(dphi, temperature, calibration: dict):
    """ Calculate pO2 in hPa for a given dphi, temperature and calibration.

    :params dphi: phase angle in °
    :param temperature: in °C
    :param calibration: dictionary as generated by the parser functions (metadata['calibration']).

    Expected entries:
        - dphi0 [°]
        - dphi100 [°]
        - temp0 [°C]
        - temp100 [°C]
        - pressure [mbar]
        - humidity [%]
        - percentO2 [%]
        - f
        - m
        - kt [1/K]
        - tt [1/K]
        - mt [1/K]
        - ft [1/K]
    """
    return calculate_pO2(dphi, temperature, **calibration)


#def calculate_pO2_from_interpolated_calibration(dphi, temperature):
#    pass
# Not Implemented yet.

####################################
# future c funcs
####################################

import ctypes
import os
import sys
if sys.platform.startswith('linux'):
    calc_oxygen = ctypes.CDLL(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'oxycalc.so'))
    calc_oxygen.calc_tau0_20.restype = ctypes.c_double
    calc_oxygen.calc_ksv_20.restype = ctypes.c_double
    calc_oxygen.calc_pO2.restype = ctypes.c_double
elif sys.platform == 'win32':
    calc_oxygen = ctypes.CDLL(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'oxycalc.dll'))
    calc_oxygen.calc_tau0_20.restype = ctypes.c_double
    calc_oxygen.calc_ksv_20.restype = ctypes.c_double
    calc_oxygen.calc_pO2.restype = ctypes.c_double
else:
    print('Warning: calculating oxygen values from dphi is not supported on this operating system.', file=sys.stderr)

def _calc_tau0_20(tau0: float, temperature: float, tt: float):
    """ Calculate tau0 for 20°C

    :param tau0: tau0 in µs
    :param temperature: in °C
    :param tt: sensor constant (1in 1/K)
    :return: tau0 at 20°C
    """
    return calc_oxygen.calc_tau0_20(ctypes.c_double(tau0), ctypes.c_double(temperature), ctypes.c_double(tt))

def _calc_ksv_20(tau0_20: float, tau100: float, pO2: float, temp100: float, f: float, m: float, kt: float, tt: float,
                 mt: float, ft: float):
    """ Calculate ksv at 20°C.

    :param tau0_20: tau at 0% O2 and 20°C
    :param tau100: tau at T and pO2
    :param pO2: pO2 for tau
    :param temp0: temperature for tau0
    :param temp100: temperature for ksv
    :param f: sensor constant
    :param m: sensor constant
    :param kt: sensor constant (in 1/K)
    :param tt: sensor constant (in 1/K)
    :param mt: sensor constant (in 1/K)
    :param ft: sensor constant (in 1/K), hardly used
    :return: ksv value at 20°C
    """
    return calc_oxygen.calc_ksv_20(ctypes.c_double(tau0_20), ctypes.c_double(tau100), ctypes.c_double(pO2),
                                    ctypes.c_double(temp100), ctypes.c_double(f), ctypes.c_double(m),
                                    ctypes.c_double(kt), ctypes.c_double(tt), ctypes.c_double(mt), ctypes.c_double(ft))


def _calc_pO2(tau: Iterable[float] | float, temperature: Iterable[float] | float, tau0_20: Iterable[float] | float,
              ksv_20: Iterable[float] | float, f: float, m: float, kt: float, tt: float, mt: float, ft: float) -> pd.Series | float:
    """ Calculate pO2 in hPa

    :param tau: tau in µs (iterable or float)
    :param temperature: sensor temperature in °C (iterable or float)
    :param tau0_20: tau0 at 20°C (iterable or float)
    :param ksv_20: ksv at 20°C (iterable or float)
    :param f: sensor constant
    :param m: sensor constant
    :param kt: sensor constant (in 1/K)
    :param tt: sensor constant (in 1/K)
    :param mt: sensor constant (in 1/K)
    :param ft: sensor constant (in 1/K), hardly used
    :return: pO2 in hPa as pd.Series of float
    """
    if isinstance(tau, (int, float)):
        tau = [tau]
    if isinstance(temperature, (int, float)):
        temperature = [temperature]
    if isinstance(tau0_20, (int, float)):
        tau0_20 = [tau0_20]
    if isinstance(ksv_20, (int, float)):
        ksv_20 = [ksv_20]

    # get the index if there is any
    if isinstance(tau, pd.Series):
        index = tau.index
    elif isinstance(temperature, pd.Series):
        index = temperature.index
    elif isinstance(tau0_20, pd.Series):
        index = tau0_20.index
    elif isinstance(ksv_20, pd.Series):
        index = ksv_20.index
    else:
        index = None

    # check if length is one or is the same
    maxlen = max(len(tau), len(temperature), len(tau0_20), len(ksv_20))
    if (len(tau) not in (1, maxlen) or len(temperature) not in (1, maxlen) or
            len(tau0_20) not in (1, maxlen) or len(ksv_20) not in (1, maxlen)):
        raise ValueError(f'operands could not be broadcast together with shapes ({len(tau)}) ({len(temperature)}) '
                         f'({len(tau0_20)}) ({len(ksv_20)})')

    # get to same length
    tau *= int(maxlen/len(tau))
    temperature *= int(maxlen/len(temperature))
    tau0_20 *= int(maxlen/len(tau0_20))
    ksv_20 *= int(maxlen/len(ksv_20))

    results = []
    for tau_, temperature_, tau0_20_, ksv_20_ in zip(tau, temperature, tau0_20, ksv_20):
        results.append(calc_oxygen.calc_pO2(ctypes.c_double(tau_), ctypes.c_double(temperature_), ctypes.c_double(tau0_20_),
                                ctypes.c_double(ksv_20_), ctypes.c_double(f), ctypes.c_double(m), ctypes.c_double(kt),
                                ctypes.c_double(tt), ctypes.c_double(mt), ctypes.c_double(ft)))
    if maxlen == 1 and index is None:
        return results[0]
    return pd.Series(results, index=index)

