"""
SB Fabric Pipeline — Meta Sound Block Tests

Tests the complete encode→modulate→readback→enrich→extract pipeline,
all running through GPU-resident Fabric primitives.

Run:
    cd afs-engine
    python3 tests/test_sb_fabric_pipeline.py -v
"""

import unittest
import numpy as np
import time
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from afs.gpu.sb_fabric import SoundBlockFabric, gpu_available
from afs.gpu.sb_fabric_pipeline import SBFabricPipeline

# Reference decision (same as Colony Sound Block tests)
SYNTH_DECISION = {
    "factors": [
        {"name": "price", "range": [2000, 6000]},
        {"name": "analog_warmth", "range": [0, 100]},
        {"name": "versatility", "range": [0, 100]},
        {"name": "build_quality", "range": [0, 100]},
        {"name": "resale_value", "range": [0, 100]},
    ],
    "relationships": [
        {"from": "price", "to": "build_quality", "strength": 0.8},
        {"from": "analog_warmth", "to": "versatility", "strength": -0.6},
        {"from": "build_quality", "to": "resale_value", "strength": 0.7},
        {"from": "versatility", "to": "analog_warmth", "strength": -0.3},
    ],
    "options": [
        {"name": "Prophet 6", "price": 3500, "analog_warmth": 85, "versatility": 45,
         "build_quality": 80, "resale_value": 75},
        {"name": "Nord Stage 4", "price": 4000, "analog_warmth": 30, "versatility": 95,
         "build_quality": 85, "resale_value": 60},
    ],
}

SIMPLE_3F = {
    "factors": [
        {"name": "speed", "range": [0, 100]},
        {"name": "cost", "range": [0, 100]},
        {"name": "quality", "range": [0, 100]},
    ],
    "relationships": [
        {"from": "speed", "to": "cost", "strength": 0.6},
        {"from": "cost", "to": "quality", "strength": -0.5},
    ],
    "options": [
        {"name": "A", "speed": 90, "cost": 30, "quality": 70},
        {"name": "B", "speed": 40, "cost": 80, "quality": 90},
    ],
}


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricPipelineBasic(unittest.TestCase):
    """Basic pipeline functionality."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)
        self.pipeline = SBFabricPipeline(self.fabric)

    def test_simple_3factor(self):
        """3-factor pipeline produces concept weights."""
        result = self.pipeline.run(
            SIMPLE_3F["factors"],
            SIMPLE_3F["relationships"],
            SIMPLE_3F["options"],
        )
        self.assertIn('concept_weights', result)
        self.assertEqual(len(result['concept_weights']), 3)
        for fn, w in result['concept_weights'].items():
            self.assertGreaterEqual(w, 0.0)
            self.assertLessEqual(w, 1.0)
        print(f"\n  3-factor weights: {result['concept_weights']}")
        print(f"  Total: {result['total_ms']:.1f}ms")

    def test_synth_decision(self):
        """5-factor synth decision pipeline."""
        result = self.pipeline.run(
            SYNTH_DECISION["factors"],
            SYNTH_DECISION["relationships"],
            SYNTH_DECISION["options"],
        )
        self.assertEqual(len(result['concept_weights']), 5)
        print(f"\n  5-factor weights: {result['concept_weights']}")
        print(f"  Relationships found: {len(result['relationships'])}")
        print(f"  Total: {result['total_ms']:.1f}ms")

    def test_with_input_modulation(self):
        """Pipeline with explicit input values for modulation."""
        result = self.pipeline.run(
            SIMPLE_3F["factors"],
            SIMPLE_3F["relationships"],
            SIMPLE_3F["options"],
            input_values={"speed": 80, "cost": 50, "quality": 70},
        )
        self.assertEqual(len(result['concept_weights']), 3)
        print(f"\n  Modulated weights: {result['concept_weights']}")

    def test_enrichment_improves_signal(self):
        """Enrichment should smooth the magnitude spectrum."""
        result_raw = self.pipeline.run(
            SIMPLE_3F["factors"],
            SIMPLE_3F["relationships"],
            SIMPLE_3F["options"],
            enrich_passes=0,
        )
        result_enriched = self.pipeline.run(
            SIMPLE_3F["factors"],
            SIMPLE_3F["relationships"],
            SIMPLE_3F["options"],
            enrich_passes=3,
        )
        # Enriched should have lower noise variance in non-peak bins
        raw_std = float(np.std(result_raw['raw_magnitudes']))
        enr_std = float(np.std(result_enriched['enriched_magnitudes']))
        print(f"\n  Raw std: {raw_std:.4f}, Enriched std: {enr_std:.4f}")


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricPipelineStages(unittest.TestCase):
    """Individual stage verification."""

    def setUp(self):
        self.fabric = SoundBlockFabric(max_size=10240)
        self.pipeline = SBFabricPipeline(self.fabric)

    def test_stage_timing(self):
        """All stages should have timing data."""
        result = self.pipeline.run(
            SYNTH_DECISION["factors"],
            SYNTH_DECISION["relationships"],
            SYNTH_DECISION["options"],
        )
        expected_stages = ['encode', 'carrier', 'modulate', 'readback', 'enrich', 'extract']
        for stage in expected_stages:
            self.assertIn(stage, result['stage_ms'],
                f"Missing timing for stage '{stage}'")

        print(f"\n  Pipeline stage timing:")
        for stage, ms in result['stage_ms'].items():
            print(f"    {stage:>10}: {ms:.2f}ms")
        print(f"    {'TOTAL':>10}: {result['total_ms']:.2f}ms")

    def test_relationships_extracted(self):
        """Phase differences should produce relationship strengths."""
        result = self.pipeline.run(
            SYNTH_DECISION["factors"],
            SYNTH_DECISION["relationships"],
            SYNTH_DECISION["options"],
        )
        # Should find at least some relationships
        rels = result['relationships']
        self.assertGreater(len(rels), 0)
        for r in rels:
            self.assertIn('from', r)
            self.assertIn('to', r)
            self.assertIn('strength', r)
            self.assertGreaterEqual(r['strength'], -1.0)
            self.assertLessEqual(r['strength'], 1.0)

    def test_fft_round_trip(self):
        """Carrier → FFT → iFFT should preserve structure."""
        result = self.pipeline.run(
            SIMPLE_3F["factors"],
            SIMPLE_3F["relationships"],
            SIMPLE_3F["options"],
            fft_size=256,
        )
        self.assertEqual(result['fft_size'], 256)
        self.assertEqual(result['n_concepts'], 3)


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestFabricPipelineBenchmark(unittest.TestCase):
    """Performance benchmark."""

    def test_pipeline_speed(self):
        """Complete pipeline should be <10ms."""
        fabric = SoundBlockFabric(max_size=10240)
        pipeline = SBFabricPipeline(fabric)

        # Warmup
        pipeline.run(SYNTH_DECISION["factors"], SYNTH_DECISION["relationships"],
                     SYNTH_DECISION["options"])

        times = []
        for _ in range(10):
            t0 = time.perf_counter()
            result = pipeline.run(
                SYNTH_DECISION["factors"],
                SYNTH_DECISION["relationships"],
                SYNTH_DECISION["options"],
            )
            times.append((time.perf_counter() - t0) * 1000)

        avg = np.mean(times)
        print(f"\n  Fabric Pipeline Benchmark (5 factors, 1024 FFT):")
        print(f"    Mean: {avg:.1f}ms ± {np.std(times):.1f}ms")
        print(f"    Min: {min(times):.1f}ms, Max: {max(times):.1f}ms")
        print(f"    vs SoundBlockService.js (CPU): ~50ms")

    def test_pipeline_10factor(self):
        """10-factor pipeline performance."""
        fabric = SoundBlockFabric(max_size=10240)
        pipeline = SBFabricPipeline(fabric)

        factors = [{"name": f"f{i}", "range": [0, 100]} for i in range(10)]
        rels = [{"from": f"f{i}", "to": f"f{(i+1)%10}", "strength": 0.5 * (-1 if i%3==0 else 1)}
                for i in range(10)]
        opts = [{"name": f"O{j}", **{f"f{i}": np.random.randint(10, 90) for i in range(10)}}
                for j in range(3)]

        # Warmup
        pipeline.run(factors, rels, opts)

        times = []
        for _ in range(10):
            t0 = time.perf_counter()
            pipeline.run(factors, rels, opts, fft_size=2048)
            times.append((time.perf_counter() - t0) * 1000)

        avg = np.mean(times)
        print(f"\n  10-factor Pipeline (2048 FFT):")
        print(f"    Mean: {avg:.1f}ms ± {np.std(times):.1f}ms")


if __name__ == "__main__":
    unittest.main(verbosity=2)
