import aquila_borg as borg
import configparser
import os
#import symbolic_pofk.linear
from functools import partial
import jax
import jax.numpy as jnp
import scipy.integrate
import math
import re

# Output stream management
cons = borg.console()
myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x))

a = jnp.array([1.0])
x = jnp.copy(a)


def compute_As(cpar):
    """
    Compute As given values of sigma8
    
    Args:
        :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with wrong As
        
    Returns:
        :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with updated As
    
    """
    
    # requires BORG-CLASS
    if not hasattr(borg.cosmo, 'ClassCosmo'):
        raise ImportError(
            "BORG-CLASS is required to compute As, but is not installed.")
        
    sigma8_true = jnp.copy(cpar.sigma8)
    cpar.sigma8 = 0
    cpar.A_s = 2.3e-9
    k_max, k_per_decade = 10, 100
    extra_class = {}
    extra_class['YHe'] = '0.24'
    cosmo = borg.cosmo.ClassCosmo(cpar, k_per_decade, k_max, extra=extra_class)
    cosmo.computeSigma8()
    cos = cosmo.getCosmology()
    cpar.A_s = (sigma8_true/cos['sigma_8'])**2*cpar.A_s
    cpar.sigma8 = sigma8_true
    
    return cpar


def get_cosmopar(ini_file):
    """
    Extract cosmological parameters from an ini file
    
    Args:
        :ini_file (str): Path to the ini file
    
    Returns:
        :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters
    """
    
    config = configparser.ConfigParser()
    config.read(ini_file)
        
    cpar = borg.cosmo.CosmologicalParameters()
    cpar.default()
    cpar.fnl = float(config['cosmology']['fnl'])
    cpar.omega_k = float(config['cosmology']['omega_k'])
    cpar.omega_m = float(config['cosmology']['omega_m'])
    cpar.omega_b = float(config['cosmology']['omega_b'])
    cpar.omega_q = float(config['cosmology']['omega_q'])
    cpar.h = float(config['cosmology']['h100'])
    cpar.sigma8 = float(config['cosmology']['sigma8'])
    cpar.n_s = float(config['cosmology']['n_s'])
    cpar.w = float(config['cosmology']['w'])
    cpar.wprime = float(config['cosmology']['wprime'])
    # cpar.A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
            # cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s)
        
    cpar = compute_As(cpar)
    
    return cpar


def get_action():
    """
    Find mode which BORG is currently being run in
    
    Returns:
        :last_line (str): The mode in which BORG is being run
    """
    
    # Find the last line of the file
    with open('trace_file.txt', 'rb') as file:
        file.seek(-2, os.SEEK_END)  # Go to the second last byte in the file.
        last_line = ''
        while len(last_line) == 0:
            while file.read(1) != b'\n':  # Keep moving back until you find a newline.
                file.seek(-2, os.SEEK_CUR)
            last_line = file.readline().decode().strip('\n')
            
    # Check that the this is the line was want and find action
    cmd = "hades_python"
    assert cmd in last_line
    # Pattern matches the command followed by a word (non-space characters), captures the word after the last occurrence
    pattern = rf'(?:{re.escape(cmd)})(?=\s)(?:\s+(\S+))'
    matches = re.findall(pattern, last_line)
    last_line = matches[-1]
    myprint(f'Running BORG mode: {last_line}')
    
    return last_line


speed_of_light = 299792  # km/s

@partial(jax.jit,)
def z_cos(r_hMpc: float, Omega_m: float) -> float:
    """
    Convert distance in Mpc/h to cosmological redshift assuming LCDM

    Args:
        - r_hMpc (jnp.ndarray): Distances to tracers in Mpc/h
        - Omega_m (float): The value of Omega_m at z=0

    Returns:
        - jnp.ndarray: The redshifts.
    """
    Omega_L = 1. - Omega_m
    q0 = Omega_m/2.0 - Omega_L
    return (1.0 - jnp.sqrt(1 - 2*r_hMpc*100*(1 + q0)/speed_of_light))/(1.0 + q0)


def get_sigma_bulk(R, cpar):
    """
    Compute the bulk flow variance when the field is smoothed with a top
    hat filter of side length R.
    
    sigma^2 = H_0^2 f^2 / (2 \pi^2) \int_0^\infty dk W^2(k,R) P(k)
    
    where W(k,R) is the Fourier transform of the top hat filter
    
    Args:
        :R (float): Box length 
        :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with wrong As
        
    Returns:
        :sigma (float): Bulk flow RMS [km/s]
        
    """
    
    k_max, k_per_decade = 10, 100
    extra_class = {}
    extra_class['YHe'] = '0.24'
    cosmo = borg.cosmo.ClassCosmo(cpar, k_per_decade, k_max, extra=extra_class)
    
    # Compute integrand
    k = jnp.logspace(-4, 1, 300) # (h / Mpc)
    plin_class = cosmo.get_Pk_matter(k) # CHECK: Is this in (Mpc/h)^3?
    x = k * R
    integrand = plin_class * (3 * (jnp.sin(x) - x * jnp.cos(x)) / x ** 3) ** 2 # (Mpc/h)^3
    
    # Check units
    sigma = scipy.integrate.simpson(integrand, x=k) # (Mpc/h)^2
    H0 = 100 # h km/s/Mpc
    f = cpar.omega_m ** 0.55
    sigma = jnp.sqrt(sigma / 2) * H0 * f / math.pi # (Mpc/h) * h km/s/Mpc = km/s

    return sigma


def parse_file_to_dict(file_path):
    config_dict = {}
    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if not line or line.startswith("#"):  # Skip empty lines and comments
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip()

            # Convert the value to the appropriate type (int, float, or leave as string)
            if value.isdigit():
                value = int(value)
            else:
                try:
                    value = float(value)
                except ValueError:
                    pass  # Leave it as a string if it cannot be converted

            config_dict[key] = value

    return config_dict


if __name__ == "__main__":
    
    cpar = get_cosmopar('../conf/basic_ini.ini')
    print(cpar)
    cpar = compute_As(cpar)
    print(cpar)
    print(1.e-9 * symbolic_pofk.linear.sigma8_to_As(
            cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s))
    
    print(get_sigma_bulk(500, cpar))
    
    
"""
- Find where else symbolic As is used
- Add prior for sigma_bulk
"""
