import numpy as np
import os
import pytest
from rvspecfit import fitter_ccf
from rvspecfit import make_ccf
from rvspecfit import serializer
from rvspecfit import spec_fit

def create_fake_ccf_data(setup_name, params, prefix):
    # params: list of parameter sets (e.g. [[5000, 2, 0], [6000, 4, -1]])
    n_models = len(params)
    n_points = 100
    
    # Fake FFT data
    # Random FFTs
    n_fft = n_points // 2 + 1
    fft = np.random.normal(size=(n_models, n_fft)) + 1j * np.random.normal(size=(n_models, n_fft))
    fft2 = np.random.normal(size=(n_models, n_fft)) + 1j * np.random.normal(size=(n_models, n_fft))
    models = np.random.normal(size=(n_models, n_points))
    
    ccfconf = {
        'logl0': 0,
        'logl1': 1,
        'npoints': n_points,
        'continuum': True,
        'maxcontpts': 10,
        'splinestep': 1000
    }
    
    dHash = {}
    dHash['params'] = np.array(params)
    dHash['ccfconf'] = ccfconf
    dHash['vsinis'] = [None] * n_models
    dHash['parnames'] = ['teff', 'logg', 'feh']
    dHash['revision'] = 'test'
    
    # Save files
    ccf_info_fname = prefix + '/' + make_ccf.get_ccf_info_name(setup_name, True)
    ccf_dat_fname = prefix + '/' + make_ccf.get_ccf_dat_name(setup_name, True)
    ccf_mod_fname = prefix + '/' + make_ccf.get_ccf_mod_name(setup_name, True)
    
    serializer.save_dict_to_hdf5(ccf_info_fname, dHash)
    np.savez(ccf_dat_fname, fft=fft, fft2=fft2)
    np.save(ccf_mod_fname, models)

def test_mismatch(tmp_path):
    prefix = str(tmp_path)
        
    # Setup A: Params 1, 2
    params_A = [[5000, 4.0, 0.0], [6000, 4.0, 0.0]]
    create_fake_ccf_data('setupA', params_A, prefix=prefix)
    
    # Setup B: Params 2, 1 (Reversed order)
    params_B = [[6000, 4.0, 0.0], [5000, 4.0, 0.0]]
    create_fake_ccf_data('setupB', params_B, prefix=prefix)
    
    # Config
    config = {
        'template_lib': prefix,
        'max_vel': 500,
        'vel_step0': 5,
        'ccf_continuum_normalize': True
    }
    
    # Fake SpecData
    lam = np.linspace(4000, 5000, 100)
    spec = np.ones_like(lam)
    espec = np.ones_like(lam) * 0.1
    
    specdata = [
        spec_fit.SpecData('setupA', lam, spec, espec),
        spec_fit.SpecData('setupB', lam, spec, espec)
    ]
    
    with pytest.raises(RuntimeError, match='The parameters of the CCF templates do not match'):
        fitter_ccf.fit(specdata, config)

def test_match(tmp_path):
    prefix = str(tmp_path)
        
    # Setup A: Params 1, 2
    params_A = [[5000, 4.0, 0.0], [6000, 4.0, 0.0]]
    create_fake_ccf_data('setupA', params_A, prefix=prefix)
    
    # Setup B: Params 1, 2 (Same order)
    params_B = [[5000, 4.0, 0.0], [6000, 4.0, 0.0]]
    create_fake_ccf_data('setupB', params_B, prefix=prefix)
    
    # Config
    config = {
        'template_lib': prefix,
        'max_vel': 500,
        'vel_step0': 5,
        'ccf_continuum_normalize': True
    }
    
    # Fake SpecData
    lam = np.linspace(4000, 5000, 100)
    spec = np.ones_like(lam)
    espec = np.ones_like(lam) * 0.1
    
    specdata = [
        spec_fit.SpecData('setupA', lam, spec, espec),
        spec_fit.SpecData('setupB', lam, spec, espec)
    ]
    
    # Should succeed
    res = fitter_ccf.fit(specdata, config)
    assert res is not None

@pytest.fixture(autouse=True)
def clear_cache():
    fitter_ccf.CCFCache.ccfs.clear()
    fitter_ccf.CCFCache.ccf2s.clear()
    fitter_ccf.CCFCache.ccf_models.clear()
    fitter_ccf.CCFCache.ccf_info.clear()

