import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import patch

from scripts.web import model


class ModelCliTests(unittest.TestCase):
    def test_map_snapshot_is_converted_to_native_sra(self):
        import numpy as np
        from scripts.web.terrain import write_map_sra
        with tempfile.TemporaryDirectory() as directory:
            root=Path(directory); raw=root/'height.f32'; np.arange(2048,dtype='<f4').tofile(raw)
            mask=root/'landmask.f32'; np.zeros(2048,dtype='<f4').tofile(mask)
            output=Path(write_map_sra(raw,'T21','height',root/'height.sra',9.80665,mask))
            lines=output.read_text().splitlines()
            self.assertEqual(int(lines[0][:10]),129)
            values=np.fromstring(lines[1],sep=' ')
            self.assertEqual(values.size,2048)
            self.assertEqual(values[0],0.0)
            self.assertEqual(values[-1],0.0)
            np.ones(2048,dtype='<f4').tofile(mask)
            output=Path(write_map_sra(raw,'T21','height',root/'height-land.sra',9.80665,mask))
            values=np.fromstring(output.read_text().splitlines()[1],sep=' ')
            self.assertAlmostEqual(values[-1],2047.0*9.80665,places=2)
            baseline=root/'baseline-land.sra'
            baseline.write_text(' '.join(['172','0','0','0','64','32','0','0'])+chr(10)+' '.join(['1']*2048))
            ocean=root/'ocean.f32';np.full(2048,100.0,dtype='<f4').tofile(ocean)
            output=Path(write_map_sra(ocean,'T21','height',root/'height-baseline.sra',9.80665,baseline))
            values=np.fromstring(output.read_text().splitlines()[1],sep=' ')
            self.assertTrue(np.all(values>0.0))
            roughness=root/'roughness.f32'; np.full(2048,2.0,dtype='<f4').tofile(roughness)
            out=Path(write_map_sra(roughness,'T21','roughness',root/'roughness.sra',9.8))
            self.assertEqual(int(out.read_text().splitlines()[0][:10]),173)
            albedo=root/'albedo.f32'; np.full(2048,.22,dtype='<f4').tofile(albedo)
            out=Path(write_map_sra(albedo,'T21','albedo',root/'albedo.sra',9.8))
            self.assertEqual(len(out.read_text().splitlines()),28)
            self.assertEqual([int(line[:10]) for line in out.read_text().splitlines()[::2]],[174]*14)
            np.full(2048,101,dtype='<f4').tofile(roughness)
            with self.assertRaises(ValueError): write_map_sra(roughness,'T21','roughness',root/'bad.sra',9.8)

    def test_full_year_steps_remain_null(self):
        with tempfile.TemporaryDirectory() as directory:
            root=Path(directory)
            runtime=root/'runtime'/'python'
            runtime.mkdir(parents=True)
            (runtime/'exoplasim').mkdir()
            config=root/'config.json'
            config.write_text('{}')
            job=root/'job.json'
            job.write_text(json.dumps({'id':'b'*32,'directory':str(root),'spec':{
                'model':{'resolution':'T21','layers':10,'ncpus':1},
                'parameters':{'noutput':24},'run':{'cycles':0,'steps':None}}}))
            calls=[]
            class FakeModel:
                def __init__(self,**values): self.workdir=values['workdir']; self.resources=[]
                def configure(self,snapshots=900,**values):
                    calls.extend([('configure',values),('snapshots',snapshots)])
                    for path in Path(self.workdir).glob('*.sra'): path.unlink()
                def cfgpostprocessor(self,**values): calls.append(('postprocessor',values))
                def _edit_namelist(self,*values): calls.append(('edit',values))
                def exportcfg(self): pass
                def run(self,*args,**values): pass
            cfg=types.ModuleType('config')
            cfg.load_config=lambda path:{'_paths':{'runtime':str(root/'runtime')},'max_result_bytes':1000}
            exoplasim=types.ModuleType('exoplasim')
            exoplasim.Model=FakeModel
            results=types.ModuleType('results')
            results.prepare_results=lambda *args,**kwargs:None
            with patch.dict(sys.modules,{'config':cfg,'exoplasim':exoplasim,'results':results}), patch.object(sys,'argv',['model','--config',str(config),'--job',str(job)]):
                model.main()
            self.assertIn(('configure',{'noutput':24,'runsteps':None}),calls)
            self.assertFalse(any(item[0]=='edit' for item in calls))
            self.assertFalse(any(item[0]=='postprocessor' for item in calls))
            for snapshots in (None,0,120,900):
                with self.subTest(snapshots=snapshots):
                    data=json.loads(job.read_text());data['spec']['run']['steps']=240
                    data['spec']['parameters']={'snapshots':snapshots} if snapshots is not None else {}
                    job.write_text(json.dumps(data));calls.clear()
                    with patch.dict(sys.modules,{'config':cfg,'exoplasim':exoplasim,'results':results}), patch.object(sys,'argv',['model','--config',str(config),'--job',str(job)]):
                        model.main()
                    self.assertIn(('snapshots',240 if snapshots is None else min(snapshots,240)),calls)
                    self.assertIn(('postprocessor',{'times':None,'timeaverage':False}),calls)
                    self.assertEqual(len([item for item in calls if item[0]=='edit']),3)
            import numpy as np
            from scripts.web import terrain,mapcatalog
            data=json.loads(job.read_text());data['spec']['parameters']={'gravity':9.8}
            data['map_snapshots']={}
            for field in mapcatalog.FIELDS:
                raw=root/f'{field}.f32';np.full(2048,mapcatalog.FIELDS[field]['default'],dtype='<f4').tofile(raw)
                data['map_snapshots'][field]=str(raw)
            job.write_text(json.dumps(data))
            with patch.dict(sys.modules,{'config':cfg,'exoplasim':exoplasim,'results':results,'terrain':terrain,'mapcatalog':mapcatalog}), patch.object(sys,'argv',['model','--config',str(config),'--job',str(job)]):
                model.main()
            for field in ('roughness','albedo'):
                native=Path(root/'runtime'/f"N032_surf_{mapcatalog.FIELDS[field]['native_code']:04d}.sra")
                self.assertEqual(native.read_bytes(),(root/'inputs/native'/f'{field}.sra').read_bytes())


if __name__=='__main__':
    unittest.main()