use std::collections::BTreeMap;
use cord_trig::{TrigGraph, TrigOp};
use crate::config::{CordicConfig, CordicOp};
use crate::lut::CordicTable;
use crate::ops::*;

/// A compiled CORDIC program ready for binary serialization or execution.
///
/// Each instruction produces its result into a slot matching its index.
/// The instruction list parallels the TrigGraph node list — compilation
/// is a direct 1:1 mapping with constants folded to fixed-point.
///
/// Supports per-operation iteration counts via `tables`: each unique
/// iteration count gets its own atan table and gain constant. Instructions
/// index into `tables` via `instr_table_idx`.
#[derive(Debug, Clone)]
pub struct CORDICProgram {
    pub word_bits: u8,
    pub instructions: Vec<CORDICInstr>,
    pub output: u32,
    /// One entry per unique iteration count (Arc-deduplicated via LutRegistry).
    pub tables: Vec<CordicTable>,
    /// Maps instruction index → table index. Only meaningful for CORDIC ops;
    /// non-CORDIC instructions have 0 here (ignored at eval time).
    pub instr_table_idx: Vec<u8>,
    /// Legacy fields for backward compat with single-table programs.
    pub atan_table: Vec<i64>,
    pub gain: i64,
}

pub struct CompileConfig {
    pub word_bits: u8,
    pub cordic: Option<CordicConfig>,
}

impl Default for CompileConfig {
    fn default() -> Self {
        Self { word_bits: 32, cordic: None }
    }
}

impl CORDICProgram {
    /// Compile a TrigGraph into a CORDIC program.
    ///
    /// Each TrigOp node becomes one CORDICInstr. Constants are converted
    /// from f64 to fixed-point at compile time. Per-operation iteration
    /// counts are resolved from the CordicConfig if present, otherwise
    /// all ops use word_bits as their iteration count.
    pub fn compile(graph: &TrigGraph, config: &CompileConfig) -> Self {
        let frac_bits = config.word_bits / 2;
        let to_fixed = |val: f64| -> i64 {
            (val * (1i64 << frac_bits) as f64).round() as i64
        };

        let cordic_cfg = config.cordic.as_ref().map(|c| c.clone())
            .unwrap_or_else(|| {
                let mut c = CordicConfig::default();
                c.word_bits = config.word_bits;
                c.default_f = config.word_bits;
                c
            });

        // Build tables for each unique iteration count, using LutRegistry for dedup
        let unique_counts = cordic_cfg.unique_iteration_counts();
        let mut count_to_table_idx: BTreeMap<u8, u8> = BTreeMap::new();
        let mut tables = Vec::new();
        for &iters in &unique_counts {
            count_to_table_idx.insert(iters, tables.len() as u8);
            tables.push(CordicTable::get(iters, frac_bits));
        }

        let default_table_idx = count_to_table_idx
            .get(&cordic_cfg.default_f)
            .copied()
            .unwrap_or_else(|| {
                let idx = tables.len() as u8;
                tables.push(CordicTable::get(cordic_cfg.default_f, frac_bits));
                idx
            });

        let resolve_table = |cop: CordicOp| -> u8 {
            let iters = cordic_cfg.iterations_for(cop);
            count_to_table_idx.get(&iters).copied().unwrap_or(default_table_idx)
        };

        let mut instr_table_idx = Vec::with_capacity(graph.nodes.len());

        let instructions: Vec<CORDICInstr> = graph.nodes.iter().map(|op| {
            let tidx = match op {
                TrigOp::Sin(_) => resolve_table(CordicOp::Sin),
                TrigOp::Cos(_) => resolve_table(CordicOp::Cos),
                TrigOp::Tan(_) => resolve_table(CordicOp::Tan),
                TrigOp::Asin(_) => resolve_table(CordicOp::Asin),
                TrigOp::Acos(_) => resolve_table(CordicOp::Acos),
                TrigOp::Atan(_) => resolve_table(CordicOp::Atan),
                TrigOp::Sinh(_) => resolve_table(CordicOp::Sinh),
                TrigOp::Cosh(_) => resolve_table(CordicOp::Cosh),
                TrigOp::Tanh(_) => resolve_table(CordicOp::Tanh),
                TrigOp::Asinh(_) => resolve_table(CordicOp::Asinh),
                TrigOp::Acosh(_) => resolve_table(CordicOp::Acosh),
                TrigOp::Atanh(_) => resolve_table(CordicOp::Atanh),
                TrigOp::Sqrt(_) => resolve_table(CordicOp::Sqrt),
                TrigOp::Exp(_) => resolve_table(CordicOp::Exp),
                TrigOp::Ln(_) => resolve_table(CordicOp::Ln),
                TrigOp::Hypot(_, _) => resolve_table(CordicOp::Hypot),
                TrigOp::Atan2(_, _) => resolve_table(CordicOp::Atan2),
                _ => 0,
            };
            instr_table_idx.push(tidx);

            match op {
                TrigOp::InputX => CORDICInstr::InputX,
                TrigOp::InputY => CORDICInstr::InputY,
                TrigOp::InputZ => CORDICInstr::InputZ,
                TrigOp::Const(v) => CORDICInstr::LoadImm(to_fixed(*v)),
                TrigOp::Add(a, b) => CORDICInstr::Add(*a, *b),
                TrigOp::Sub(a, b) => CORDICInstr::Sub(*a, *b),
                TrigOp::Mul(a, b) => CORDICInstr::Mul(*a, *b),
                TrigOp::Div(a, b) => CORDICInstr::Div(*a, *b),
                TrigOp::Neg(a) => CORDICInstr::Neg(*a),
                TrigOp::Abs(a) => CORDICInstr::Abs(*a),
                TrigOp::Sin(a) => CORDICInstr::Sin(*a),
                TrigOp::Cos(a) => CORDICInstr::Cos(*a),
                TrigOp::Tan(a) => CORDICInstr::Tan(*a),
                TrigOp::Asin(a) => CORDICInstr::Asin(*a),
                TrigOp::Acos(a) => CORDICInstr::Acos(*a),
                TrigOp::Atan(a) => CORDICInstr::Atan(*a),
                TrigOp::Sinh(a) => CORDICInstr::Sinh(*a),
                TrigOp::Cosh(a) => CORDICInstr::Cosh(*a),
                TrigOp::Tanh(a) => CORDICInstr::Tanh(*a),
                TrigOp::Asinh(a) => CORDICInstr::Asinh(*a),
                TrigOp::Acosh(a) => CORDICInstr::Acosh(*a),
                TrigOp::Atanh(a) => CORDICInstr::Atanh(*a),
                TrigOp::Sqrt(a) => CORDICInstr::Sqrt(*a),
                TrigOp::Exp(a) => CORDICInstr::Exp(*a),
                TrigOp::Ln(a) => CORDICInstr::Ln(*a),
                TrigOp::Hypot(a, b) => CORDICInstr::Hypot(*a, *b),
                TrigOp::Atan2(a, b) => CORDICInstr::Atan2(*a, *b),
                TrigOp::Min(a, b) => CORDICInstr::Min(*a, *b),
                TrigOp::Max(a, b) => CORDICInstr::Max(*a, *b),
                TrigOp::Clamp { val, lo, hi } => CORDICInstr::Clamp {
                    val: *val,
                    lo: *lo,
                    hi: *hi,
                },
            }
        }).collect();

        let primary = tables.first().cloned()
            .unwrap_or_else(|| CordicTable::for_word_bits(config.word_bits));

        CORDICProgram {
            word_bits: config.word_bits,
            instructions,
            output: graph.output,
            tables,
            instr_table_idx,
            atan_table: primary.atan.to_vec(),
            gain: primary.gain,
        }
    }

    /// Serialize to binary representation.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::new();

        // Header: magic + word_bits
        buf.extend_from_slice(b"CORD");
        buf.push(self.word_bits);

        // Atan table
        let table_len = self.atan_table.len() as u16;
        buf.extend_from_slice(&table_len.to_le_bytes());
        for &val in &self.atan_table {
            buf.extend_from_slice(&val.to_le_bytes());
        }

        // Gain
        buf.extend_from_slice(&self.gain.to_le_bytes());

        // Output slot
        buf.extend_from_slice(&self.output.to_le_bytes());

        // Instructions
        let instr_len = self.instructions.len() as u32;
        buf.extend_from_slice(&instr_len.to_le_bytes());
        for instr in &self.instructions {
            encode_instruction(&mut buf, instr);
        }

        buf
    }

    /// Deserialize from binary.
    pub fn from_bytes(data: &[u8]) -> Option<Self> {
        if data.len() < 5 || &data[0..4] != b"CORD" {
            return None;
        }
        let word_bits = data[4];
        let mut pos = 5;

        // Atan table
        let table_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
        pos += 2;
        let mut atan_table = Vec::with_capacity(table_len);
        for _ in 0..table_len {
            let val = i64::from_le_bytes(data[pos..pos + 8].try_into().ok()?);
            atan_table.push(val);
            pos += 8;
        }

        // Gain
        let gain = i64::from_le_bytes(data[pos..pos + 8].try_into().ok()?);
        pos += 8;

        // Output slot
        let output = u32::from_le_bytes(data[pos..pos + 4].try_into().ok()?);
        pos += 4;

        // Instructions
        let instr_len = u32::from_le_bytes(data[pos..pos + 4].try_into().ok()?) as usize;
        pos += 4;
        let mut instructions = Vec::with_capacity(instr_len);
        for _ in 0..instr_len {
            let (instr, consumed) = decode_instruction(&data[pos..])?;
            instructions.push(instr);
            pos += consumed;
        }

        let tables = vec![CordicTable::get(word_bits, word_bits - 1)];
        let instr_table_idx = vec![0u8; instructions.len()];

        Some(CORDICProgram {
            word_bits, instructions, output, tables, instr_table_idx, atan_table, gain,
        })
    }
}
