import pytest
import ionbench
from ionbench import modification
import numpy as np

class Test_Modifications:
    bm = ionbench.problems.staircase.HH_Benchmarker()
    def test_bounds(self):
        #Check bounds are correctly applied for all settings
        #No bounds initially
        assert not self.bm._bounded
        #Positive
        mod = modification.Modification(bounds = 'positive')
        mod.apply(self.bm)
        assert self.bm._bounded
        assert all(np.array(self.bm.lb)==0)
        assert all(np.array(self.bm.ub)==np.inf)
        #Sampler
        mod = modification.Modification(bounds = 'Sampler')
        mod.apply(self.bm)
        assert self.bm._bounded
        assert all(np.array(self.bm.lb)>0)
        assert all(np.array(self.bm.ub)<np.inf)
        #Custom
        lb = np.random.rand(self.bm.n_parameters())
        ub = np.random.rand(self.bm.n_parameters())+1
        mod = modification.Modification(bounds = 'Custom', customBounds = [lb,ub])
        mod.apply(self.bm)
        assert self.bm._bounded
        assert all(self.bm.lb==lb)
        assert all(self.bm.ub==ub)
        #None
        mod = modification.Modification()
        mod.apply(self.bm)
        assert not self.bm._bounded
    
    def test_log_transform(self):
        #Check log transform is correctly applied for all settings
        #No log transform initially
        assert not any(self.bm._logTransformParams)
        #Standard log transform
        mod = modification.Modification(logTransform = 'standard')
        mod.apply(self.bm)
        assert any(self.bm._logTransformParams)
        assert not all(self.bm._logTransformParams)
        #Full log transform
        mod = modification.Modification(logTransform = 'Full')
        mod.apply(self.bm)
        assert all(self.bm._logTransformParams)
        #Custom
        transforms = list(np.random.choice([True, False],self.bm.n_parameters()-2))+[True, False] #Random true and falses in which the last two are True and False
        mod = modification.Modification(logTransform = 'custom', customLogTransform = transforms)
        mod.apply(self.bm)
        assert any(self.bm._logTransformParams)
        assert all(np.array(self.bm._logTransformParams) == np.array(transforms))
        assert not all(self.bm._logTransformParams)
        #No log transform
        mod = modification.Modification()
        mod.apply(self.bm)
        assert not any(self.bm._logTransformParams)
    
    def test_scale_factors(self):
        #Check scale factors are correctly applied
        #No scale factors initially
        assert not self.bm._useScaleFactors
        #Turn on scale factors
        mod = modification.Modification(scaleFactors = 'on')
        mod.apply(self.bm)
        assert self.bm._useScaleFactors
        #Default is off
        mod = modification.Modification()
        mod.apply(self.bm)
        assert not self.bm._useScaleFactors
    
    def test_multiple_settings(self):
        #Check multiple settings applied at once doesn't break anything and order in the .apply() method doesn't matter
        mod = modification.Modification(logTransform='Full',bounds='positive',scaleFactors='On')
        mod.apply_log_transforms(mod.dict['log transform'], self.bm)
        mod.apply_scale_factors(mod.dict['scale factors'], self.bm)
        mod.apply_bounds(mod.dict['bounds'], self.bm)
        assert all(self.bm._logTransformParams)
        assert self.bm._useScaleFactors
        assert self.bm._bounded
        lb = self.bm.lb
        ub = self.bm.ub
        #Check order of applying doesnt change bounds
        emptyMod = modification.Empty()
        emptyMod.apply(self.bm)
        mod = modification.Modification(logTransform='Full',bounds='positive',scaleFactors='on')
        mod.apply_bounds(mod.dict['bounds'], self.bm)
        mod.apply_log_transforms(mod.dict['log transform'], self.bm)
        mod.apply_scale_factors(mod.dict['scale factors'], self.bm)
        assert all(self.bm._logTransformParams)
        assert self.bm._useScaleFactors
        assert self.bm._bounded
        assert all(np.array(self.bm.lb) == np.array(lb))
        assert all(np.array(self.bm.ub) == np.array(ub))
    
    def test_other_problems(self):
        #Basic check for another benchmarker
        mod = modification.Modification(logTransform = 'standard', bounds = 'sampler', scaleFactors = 'On')
        newbm = ionbench.problems.loewe2016.ikr()
        mod.apply(newbm)
        assert newbm._useScaleFactors
        assert all(np.array(newbm._logTransformParams) == np.array(newbm.standardLogTransform))
        assert any(newbm._logTransformParams)
        assert not all(newbm._logTransformParams)
        assert newbm._bounded
