from fractions import Fraction
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import Mock,patch

import numpy as np

from model_tooling.fire_video import parser,validate
from model_tooling.video import VideoError
from model_tooling.fire_video_render import select_candidates,views,render,WIDTH,HEIGHT,HEADER,NAMES
from model_tooling.fire_video_worker import Encoder,selected_frames,budget,distribution


def config():
    return dict(mode='fp16',excluded='off',fps=10,classes=list(NAMES),
                thresholds=dict(zip(NAMES,(.25,.25))),excluded_limit=20)


class FireVideoTests(unittest.TestCase):
    def test_options_and_modes(self):
        args=parser().parse_args(['--input','video']);validate(args)
        self.assertEqual(args.mode,'fp16');self.assertEqual(args.excluded,'off')
        for flag,val in [('--fps','.5'),('--fps','nan'),('--smoke-threshold','-1'),('--flame-threshold','inf'),
                         ('--end','0'),('--excluded-limit','101'),('--max-frames','0')]:
            with self.subTest(flag=flag),self.assertRaises(VideoError):validate(parser().parse_args(['--input','video',flag,val]))
        with self.assertRaises(VideoError):views('compare','compare')
        self.assertEqual(views('reference','compare'),[('reference','off'),('reference','on')])

    def test_threshold_class_validity_and_ties(self):
        rows=[[0,.25,10,10,20,20],[1,.249,10,10,20,20]]+[[0,.1,1,1,4,4] for _ in range(25)]
        shown,ghosts,counts=select_candidates(rows,dict(zip(NAMES,(.25,.25))),list(NAMES))
        self.assertEqual(len(shown),1);self.assertEqual(len(ghosts),20)
        self.assertEqual([i for i,_ in ghosts],list(range(1,21)))
        self.assertEqual(counts['excluded_omitted'],6)
        selected,hidden,counts=select_candidates([[2,.9,1,1,1,1],[0,.9,0,0,0,1],[1,.9,1,1,1,1]],
                                               dict(zip(NAMES,(.25,.25))),['smoke_visible'])
        self.assertFalse(selected);self.assertFalse(hidden)
        self.assertEqual(counts['invalid'],2);self.assertEqual(counts['unselected'],1)

    def test_display_only_and_rgb(self):
        cfg=config();row=dict(source_frame=0,pts_s=0,models={'fp16':{'detections':[[0,.1,10,10,80,60],[1,.9,100,50,50,50]]}})
        original=json.dumps(row,sort_keys=True)
        rgb=np.full((180,320,3),(190,30,15),dtype=np.uint8)
        off,_=render(rgb,row,cfg)
        on,counts=render(rgb,row,dict(cfg,excluded='on'))
        both,_=render(rgb,row,dict(cfg,excluded='compare'))
        self.assertEqual(both.size,(WIDTH*2,HEIGHT+HEADER))
        self.assertTrue(np.array_equal(np.asarray(both)[:,:WIDTH],np.asarray(off)))
        self.assertTrue(np.array_equal(np.asarray(both)[:,WIDTH:],np.asarray(on)))
        self.assertEqual(off.getpixel((WIDTH-5,HEADER+HEIGHT-5)),(190,30,15))
        self.assertEqual(counts['fp16:on']['excluded_drawn'],1)
        self.assertEqual(json.dumps(row,sort_keys=True),original)

    def test_pts_encode_selection_and_no_duplicate(self):
        from PIL import Image
        import av
        with tempfile.TemporaryDirectory() as root:
            path=Path(root)/'fixture.mp4'
            encoder=Encoder(path,10,WIDTH)
            for t in (Fraction(0),Fraction(1,10),Fraction(1,5)):
                encoder.write(Image.new('RGB',(WIDTH,HEIGHT+HEADER),(190,30,15)),t)
            encoder.finish(Fraction(7,30))
            frames=list(selected_frames(path,0,.24,10))
            self.assertEqual(len(frames),3)
            self.assertEqual([v[1] for v in frames],[Fraction(0),Fraction(1,10),Fraction(1,5)])
            with av.open(path) as f:self.assertAlmostEqual(float(f.duration)/1000000,7/30,places=3)
            encoder=Encoder(Path(root)/'duplicate.mp4',10,WIDTH)
            try:
                encoder.write(Image.new('RGB',(WIDTH,HEIGHT+HEADER)),Fraction(0))
                with self.assertRaises(VideoError):encoder.write(Image.new('RGB',(WIDTH,HEIGHT+HEADER)),Fraction(0))
            finally:encoder.abort()

    def test_resource_budget_and_model_close(self):
        with tempfile.TemporaryDirectory() as root,patch('model_tooling.fire_video_worker.OUT',Path(root)):
            with patch('shutil.disk_usage',return_value=Mock(free=0)),self.assertRaises(VideoError):budget({'max_output_bytes':10})
        from model_tooling.fire_video_models import Models
        obj=Models.__new__(Models);obj.handle=12;obj.lib=Mock();obj.reference=Mock();obj.torch=Mock()
        obj.close();obj.close()
        obj.lib.close_probe.assert_called_once_with(12)
        self.assertIsNone(obj.reference)
        stats=distribution([10,20,30]);self.assertEqual(stats['fps'],50)


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