import json
import io
import pathlib
import struct
import tempfile
import unittest
from unittest.mock import patch

import numpy as np

from pilot.analyze import (build_blocks, exact_pair, object_record, read_object,
                           sample_object, sharing_groups)
from pilot.tensors import decode, load_tensors, metrics


class PilotTests(unittest.TestCase):
    def test_directional_replacement_error_and_scaling(self):
        x = np.array([1., 2., 3.])
        m = metrics(x, 2 * x)
        self.assertAlmostEqual(m['cosine'], 1.)
        self.assertAlmostEqual(m['error_member_to_rep'], 1.)
        self.assertAlmostEqual(m['error_rep_to_member'], .5)
        self.assertAlmostEqual(m['scaled_error'], 0.)
        self.assertAlmostEqual(m['scale'], .5)

    def test_bfloat16(self):
        raw = np.array([0x3f80, 0xc000, 0x0000], dtype='<u2')
        np.testing.assert_array_equal(decode(raw, 'BF16'), [1., -2., 0.])

    def test_ranged_download_preserves_prefix(self):
        from pilot.download import ranged_download
        data = b'the entire known checkpoint contents'
        class Response(io.BytesIO):
            status = 206
        def request(req, timeout):
            span = req.headers['Range'].split('=')[1]
            a, b = map(int, span.split('-'))
            response = Response(data[a:b + 1])
            response.headers = {'Content-Range': f'bytes {a}-{b}/{len(data)}'}
            return response
        with tempfile.TemporaryDirectory() as td:
            part = pathlib.Path(td) / 'file.partial'
            part.write_bytes(data[:7])
            with patch('pilot.download.urllib.request.urlopen', side_effect=request):
                ranged_download('https://example.invalid/file', part, len(data), workers=2)
            self.assertEqual(part.read_bytes(), data)
            self.assertFalse(part.with_suffix('.partial.ranges.json').exists())

    def test_no_transitive_chaining(self):
        objects = [{'id': str(i)} for i in range(3)]
        pairs = [{'a': 0, 'b': 1, 'error_member_to_rep': .01, 'error_rep_to_member': .5},
                 {'a': 1, 'b': 2, 'error_member_to_rep': .01, 'error_rep_to_member': .5}]
        groups = sharing_groups(objects, pairs, .02)
        self.assertFalse(any(set(g['members']) == {0, 1, 2} for g in groups))
        self.assertTrue(all(g['max_error'] <= .02 for g in groups))

    def test_3d_reader_and_exact_block(self):
        x = np.arange(8 * 4 * 4, dtype='<f4').reshape(8, 4, 4)
        names = ['model.layers.0.conv.weight', 'model.layers.1.conv.weight']
        header = {name: {'dtype': 'F32', 'shape': list(x.shape),
                         'data_offsets': [i * x.nbytes, (i + 1) * x.nbytes]}
                  for i, name in enumerate(names)}
        encoded = json.dumps(header).encode()
        with tempfile.TemporaryDirectory() as td:
            path = pathlib.Path(td)
            (path / 'model.safetensors').write_bytes(struct.pack('<Q', len(encoded)) + encoded + x.tobytes() * 2)
            tensors = load_tensors(path)
            a = object_record(tensors[names[0]], 'block', (4, 2, 2), (4, 2, 2), '3d')
            b = object_record(tensors[names[1]], 'block', (4, 2, 2), (4, 2, 2), '3d')
            np.testing.assert_array_equal(read_object(a, tensors), x[4:8, 2:4, 2:4].ravel())
            np.testing.assert_array_equal(sample_object(a, tensors), read_object(a, tensors))
            self.assertEqual(exact_pair(a, b, tensors, float('inf'))['error_member_to_rep'], 0.)
            blocks, coverage, totals = build_blocks(tensors, 1, float('inf'), 2)
            self.assertTrue(blocks)
            self.assertTrue(all(r['sampled_blocks'] <= 2 for r in coverage))
            self.assertTrue(all(r['bytes'] + r['edge_bytes'] == x.nbytes * 2 for r in totals.values()))

    def test_end_to_end_synthetic_groups_and_savings(self):
        from pilot.analyze import main
        arrays = {}
        for layer in (0, 1):
            factor = 1 + layer * .005
            arrays[f'model.layers.{layer}.proj.weight'] = (np.arange(64 * 64, dtype='<f4').reshape(64, 64) + 1) * factor
            arrays[f'model.layers.{layer}.conv.weight'] = (np.arange(8 * 4 * 4, dtype='<f4').reshape(8, 4, 4) + 1) * factor
        header, payload, offset = {}, [], 0
        for name, x in arrays.items():
            header[name] = {'dtype': 'F32', 'shape': list(x.shape), 'data_offsets': [offset, offset + x.nbytes]}
            payload.append(x.tobytes())
            offset += x.nbytes
        encoded = json.dumps(header).encode()
        with tempfile.TemporaryDirectory() as td:
            root = pathlib.Path(td)
            (root / 'weights').mkdir()
            (root / 'weights/model.safetensors').write_bytes(struct.pack('<Q', len(encoded)) + encoded + b''.join(payload))
            (root / 'manifest.json').write_text(json.dumps({'repo': 'synthetic', 'revision': 'test'}))
            with patch('sys.argv', ['analyze', '--root', str(root), '--seconds', '60']):
                main()
            summary = json.loads((root / 'analysis/summary.json').read_text())
            self.assertTrue(summary['completed'])
            groups = json.loads((root / 'analysis/groups_tensor_whole_0.01.json').read_text())
            self.assertEqual(len(groups), 2)
            self.assertTrue(all(len(g['members']) == 2 for g in groups))
            self.assertTrue(all(g['max_error'] < .01 for g in groups))
            from pilot.plot import main as plot_main
            from pilot.validate import validate
            with patch('sys.argv', ['plot', '--root', str(root)]):
                plot_main()
            self.assertTrue(validate(root)['passed'])
            self.assertEqual(len(list((root / 'figures').glob('*.pdf'))), 6)
            import csv
            from collections import defaultdict
            with (root / 'analysis/null_histograms.csv').open() as f:
                bins = list(csv.DictReader(f))
            areas = defaultdict(float)
            counts = defaultdict(int)
            sizes = {}
            for row in bins:
                key = (row['mode'], row['ndim'], row['method'])
                areas[key] += float(row['density']) * (float(row['bin_right']) - float(row['bin_left']))
                counts[key] += int(row['count'])
                sizes[key] = int(row['samples'])
            self.assertTrue(areas)
            for key in areas:
                self.assertAlmostEqual(areas[key], 1.)
                self.assertEqual(counts[key], sizes[key])
            with (root / 'analysis/heatmap_pairs.csv').open() as f:
                cells = list(csv.DictReader(f))
            self.assertTrue(all(row['verified'] == 'False' for row in cells if row['identity'] == 'True'))


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