finetype-model 0.1.0

//! Training utilities for tiered inference models.
//!
//! Trains a hierarchy of CharCNN models:
//! - Tier 0: Broad type classifier (e.g., VARCHAR, DATE, TIMESTAMP)
//! - Tier 1: Per-broad-type category models (e.g., VARCHAR → internet/person/code)
//! - Tier 2: Per-category type models for categories with many types

use crate::char_cnn::{CharCnn, CharCnnConfig, CharVocab};
use candle_core::{DType, Device, Tensor};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use finetype_core::{Sample, Taxonomy, TierGraph};
use rand::seq::SliceRandom;
use std::collections::HashMap;
use std::path::Path;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum TieredTrainingError {
    #[error("Model error: {0}")]
    ModelError(#[from] candle_core::Error),
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
    #[error("No samples for tier {tier} group {group}")]
    EmptyGroup { tier: String, group: String },
}

/// Configuration for tiered training.
#[derive(Debug, Clone)]
pub struct TieredTrainingConfig {
    pub batch_size: usize,
    pub epochs: usize,
    pub learning_rate: f64,
    pub max_seq_length: usize,
    pub embed_dim: usize,
    pub num_filters: usize,
    pub hidden_dim: usize,
    pub weight_decay: f64,
    /// Minimum number of types in a (broad_type, category) group to train a Tier 2 model.
    /// Groups with fewer types are resolved by Tier 1 directly.
    pub tier2_min_types: usize,
}

impl Default for TieredTrainingConfig {
    fn default() -> Self {
        Self {
            batch_size: 64,
            epochs: 10,
            learning_rate: 1e-3,
            max_seq_length: 128,
            embed_dim: 32,
            num_filters: 64,
            hidden_dim: 128,
            weight_decay: 1e-4,
            tier2_min_types: 1,
        }
    }
}

/// Trainer for the full tiered model graph.
pub struct TieredTrainer {
    config: TieredTrainingConfig,
    device: Device,
    vocab: CharVocab,
}

impl TieredTrainer {
    /// Create a new tiered trainer.
    pub fn new(config: TieredTrainingConfig) -> Self {
        let device = Self::get_device();
        let vocab = CharVocab::new();
        Self {
            config,
            device,
            vocab,
        }
    }

    /// Train all tier models from a single dataset.
    ///
    /// Saves models to `output_dir` with the following structure:
    /// ```text
    /// output_dir/
    ///   tier0/                       # Broad type classifier
    ///   tier1_{broad_type}/          # Category classifier per broad type
    ///   tier2_{broad_type}_{category}/  # Type classifier per category
    ///   tier_graph.json              # Graph metadata
    /// ```
    pub fn train_all(
        &self,
        taxonomy: &Taxonomy,
        samples: &[Sample],
        output_dir: &Path,
    ) -> Result<TieredTrainingReport, TieredTrainingError> {
        let graph = taxonomy.tier_graph();
        let mut report = TieredTrainingReport::default();

        eprintln!("=== Tiered Training ===");
        eprintln!("{}", graph.summary());

        // --- Tier 0: Broad type classification ---
        eprintln!("\n--- Training Tier 0 (broad type) ---");
        let tier0_dir = output_dir.join("tier0");
        let tier0_accuracy = self.train_tier0(&graph, samples, &tier0_dir)?;
        report.tier0_accuracy = tier0_accuracy;
        report.tier0_classes = graph.num_broad_types();
        eprintln!(
            "Tier 0: {:.2}% accuracy ({} classes)",
            tier0_accuracy * 100.0,
            graph.num_broad_types()
        );

        // --- Tier 1: Per-broad-type category models ---
        for broad_type in graph.broad_types() {
            let categories = graph.categories_for(broad_type);
            if categories.len() <= 1 {
                // Only one category — no Tier 1 model needed, resolved by lookup
                eprintln!(
                    "\n--- Tier 1 [{}]: skipped (single category: {}) ---",
                    broad_type,
                    categories.first().map(|s| s.as_str()).unwrap_or("none")
                );
                report.tier1_skipped.push(broad_type.clone());
                continue;
            }

            eprintln!(
                "\n--- Training Tier 1 [{}] ({} categories) ---",
                broad_type,
                categories.len()
            );
            let tier1_dir = output_dir.join(format!("tier1_{}", broad_type));
            let tier1_accuracy = self.train_tier1(&graph, broad_type, samples, &tier1_dir)?;
            report.tier1_results.push(TierModelResult {
                name: broad_type.clone(),
                classes: categories.len(),
                accuracy: tier1_accuracy,
            });
            eprintln!(
                "Tier 1 [{}]: {:.2}% accuracy ({} categories)",
                broad_type,
                tier1_accuracy * 100.0,
                categories.len()
            );
        }

        // --- Tier 2: Per-category type models ---
        for broad_type in graph.broad_types() {
            for category in graph.categories_for(broad_type) {
                let n_types = graph.num_types(broad_type, category);
                if n_types <= self.config.tier2_min_types {
                    // Too few types — resolved directly by Tier 1
                    continue;
                }

                eprintln!(
                    "\n--- Training Tier 2 [{}/{}] ({} types) ---",
                    broad_type, category, n_types
                );
                let tier2_dir = output_dir.join(format!("tier2_{}_{}", broad_type, category));
                let tier2_accuracy =
                    self.train_tier2(&graph, broad_type, category, samples, &tier2_dir)?;
                report.tier2_results.push(TierModelResult {
                    name: format!("{}/{}", broad_type, category),
                    classes: n_types,
                    accuracy: tier2_accuracy,
                });
                eprintln!(
                    "Tier 2 [{}/{}]: {:.2}% accuracy ({} types)",
                    broad_type,
                    category,
                    tier2_accuracy * 100.0,
                    n_types
                );
            }
        }

        // Save tier graph metadata
        let graph_meta = self.build_graph_metadata(&graph);
        let graph_json =
            serde_json::to_string_pretty(&graph_meta).unwrap_or_else(|_| "{}".to_string());
        std::fs::create_dir_all(output_dir)?;
        std::fs::write(output_dir.join("tier_graph.json"), graph_json)?;

        eprintln!("\n=== Tiered Training Complete ===");
        eprintln!("{}", report);

        Ok(report)
    }

    /// Train Tier 0: broad type classification.
    fn train_tier0(
        &self,
        graph: &TierGraph,
        samples: &[Sample],
        output_dir: &Path,
    ) -> Result<f64, TieredTrainingError> {
        // Relabel samples: full label → broad_type
        let relabeled: Vec<Sample> = samples
            .iter()
            .filter_map(|s| {
                graph.broad_type_for(&s.label).map(|bt| Sample {
                    text: s.text.clone(),
                    label: bt.to_string(),
                })
            })
            .collect();

        if relabeled.is_empty() {
            return Err(TieredTrainingError::EmptyGroup {
                tier: "0".into(),
                group: "all".into(),
            });
        }

        // Build label set
        let mut labels: Vec<String> = graph.broad_types().to_vec();
        labels.sort();

        self.train_model(&labels, &relabeled, output_dir)
    }

    /// Train Tier 1: category classification within a broad type.
    fn train_tier1(
        &self,
        graph: &TierGraph,
        broad_type: &str,
        samples: &[Sample],
        output_dir: &Path,
    ) -> Result<f64, TieredTrainingError> {
        // Filter samples for this broad type, relabel to category
        let relabeled: Vec<Sample> = samples
            .iter()
            .filter_map(|s| {
                let (bt, cat) = graph.tier_path(&s.label)?;
                if bt == broad_type {
                    Some(Sample {
                        text: s.text.clone(),
                        label: cat.clone(),
                    })
                } else {
                    None
                }
            })
            .collect();

        if relabeled.is_empty() {
            return Err(TieredTrainingError::EmptyGroup {
                tier: "1".into(),
                group: broad_type.into(),
            });
        }

        let mut labels: Vec<String> = graph.categories_for(broad_type).to_vec();
        labels.sort();

        self.train_model(&labels, &relabeled, output_dir)
    }

    /// Train Tier 2: type classification within a (broad_type, category).
    fn train_tier2(
        &self,
        graph: &TierGraph,
        broad_type: &str,
        category: &str,
        samples: &[Sample],
        output_dir: &Path,
    ) -> Result<f64, TieredTrainingError> {
        // Filter samples for this (broad_type, category), keep original label
        let filtered: Vec<Sample> = samples
            .iter()
            .filter(|s| {
                graph
                    .tier_path(&s.label)
                    .map(|(bt, cat)| bt == broad_type && cat == category)
                    .unwrap_or(false)
            })
            .cloned()
            .collect();

        if filtered.is_empty() {
            return Err(TieredTrainingError::EmptyGroup {
                tier: "2".into(),
                group: format!("{}/{}", broad_type, category),
            });
        }

        let mut labels: Vec<String> = graph.types_for(broad_type, category).to_vec();
        labels.sort();

        self.train_model(&labels, &filtered, output_dir)
    }

    /// Train a single CharCNN model with the given labels and samples.
    /// Returns the final epoch accuracy.
    fn train_model(
        &self,
        labels: &[String],
        samples: &[Sample],
        output_dir: &Path,
    ) -> Result<f64, TieredTrainingError> {
        let n_classes = labels.len();
        let label_to_index: HashMap<String, usize> = labels
            .iter()
            .enumerate()
            .map(|(i, l)| (l.clone(), i))
            .collect();

        // Shuffle samples
        let mut samples_vec: Vec<&Sample> = samples.iter().collect();
        let mut rng = rand::thread_rng();
        samples_vec.shuffle(&mut rng);

        // Initialize model
        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &self.device);

        let model_config = CharCnnConfig {
            vocab_size: self.vocab.vocab_size(),
            max_seq_length: self.config.max_seq_length,
            embed_dim: self.config.embed_dim,
            num_filters: self.config.num_filters,
            hidden_dim: self.config.hidden_dim,
            n_classes,
            ..Default::default()
        };

        let model = CharCnn::new(model_config, vb)?;

        // Create optimizer
        let params = ParamsAdamW {
            lr: self.config.learning_rate,
            weight_decay: self.config.weight_decay,
            ..Default::default()
        };
        let mut optimizer = AdamW::new(varmap.all_vars(), params)?;

        // Training loop
        let num_batches = samples_vec.len().div_ceil(self.config.batch_size);
        let mut final_accuracy = 0.0;

        for epoch in 0..self.config.epochs {
            if epoch > 0 {
                samples_vec.shuffle(&mut rng);
            }

            let mut total_loss = 0.0;
            let mut num_correct = 0usize;
            let mut num_total = 0usize;

            for batch_idx in 0..num_batches {
                let start = batch_idx * self.config.batch_size;
                let end = (start + self.config.batch_size).min(samples_vec.len());
                let batch: Vec<&Sample> = samples_vec[start..end].to_vec();

                let (input_ids, batch_labels) = self.prepare_batch(&batch, &label_to_index)?;

                let logits = model.forward(&input_ids)?;
                let logits = logits.contiguous()?;
                let loss = candle_nn::loss::cross_entropy(&logits, &batch_labels)?;
                optimizer.backward_step(&loss)?;

                let loss_val = loss.to_scalar::<f32>()?;
                total_loss += loss_val;

                let predictions = logits.argmax(1)?;
                let correct = predictions
                    .eq(&batch_labels)?
                    .to_dtype(DType::F32)?
                    .sum_all()?
                    .to_scalar::<f32>()?;
                num_correct += correct as usize;
                num_total += batch.len();
            }

            let avg_loss = total_loss / num_batches as f32;
            final_accuracy = num_correct as f64 / num_total as f64;

            eprintln!(
                "  Epoch {}/{}: loss={:.4}, accuracy={:.2}%",
                epoch + 1,
                self.config.epochs,
                avg_loss,
                final_accuracy * 100.0
            );
        }

        // Save model
        std::fs::create_dir_all(output_dir)?;
        varmap.save(output_dir.join("model.safetensors"))?;

        // Save config
        let config_str = format!(
            "vocab_size: {}\nmax_seq_length: {}\nembed_dim: {}\nnum_filters: {}\nhidden_dim: {}\nn_classes: {}\nmodel_type: char_cnn\n",
            self.vocab.vocab_size(),
            self.config.max_seq_length,
            self.config.embed_dim,
            self.config.num_filters,
            self.config.hidden_dim,
            n_classes
        );
        std::fs::write(output_dir.join("config.yaml"), config_str)?;

        // Save label mapping
        let labels_json = serde_json::to_string_pretty(labels).unwrap_or_else(|_| "[]".to_string());
        std::fs::write(output_dir.join("labels.json"), labels_json)?;

        Ok(final_accuracy)
    }

    /// Prepare a batch for training.
    fn prepare_batch(
        &self,
        samples: &[&Sample],
        label_to_index: &HashMap<String, usize>,
    ) -> Result<(Tensor, Tensor), TieredTrainingError> {
        let batch_size = samples.len();
        let max_len = self.config.max_seq_length;

        let mut all_ids = Vec::with_capacity(batch_size * max_len);
        let mut all_labels = Vec::with_capacity(batch_size);

        for sample in samples {
            let ids = self.vocab.encode(&sample.text, max_len);
            all_ids.extend(ids);
            let label_idx = label_to_index.get(&sample.label).copied().unwrap_or(0) as u32;
            all_labels.push(label_idx);
        }

        let input_ids = Tensor::new(all_ids, &self.device)?.reshape((batch_size, max_len))?;
        let labels = Tensor::new(all_labels, &self.device)?;

        Ok((input_ids, labels))
    }

    /// Build graph metadata JSON for the inference engine.
    fn build_graph_metadata(&self, graph: &TierGraph) -> serde_json::Value {
        use serde_json::json;

        let mut tier1_models = serde_json::Map::new();
        for broad_type in graph.broad_types() {
            let categories = graph.categories_for(broad_type);
            if categories.len() > 1 {
                tier1_models.insert(
                    broad_type.clone(),
                    json!({
                        "dir": format!("tier1_{}", broad_type),
                        "categories": categories,
                    }),
                );
            } else {
                tier1_models.insert(
                    broad_type.clone(),
                    json!({
                        "direct": categories.first(),
                    }),
                );
            }
        }

        let mut tier2_models = serde_json::Map::new();
        for broad_type in graph.broad_types() {
            for category in graph.categories_for(broad_type) {
                let n_types = graph.num_types(broad_type, category);
                let key = format!("{}_{}", broad_type, category);
                if n_types > self.config.tier2_min_types {
                    tier2_models.insert(
                        key,
                        json!({
                            "dir": format!("tier2_{}_{}", broad_type, category),
                            "types": graph.types_for(broad_type, category),
                            "count": n_types,
                        }),
                    );
                } else {
                    let types = graph.types_for(broad_type, category);
                    tier2_models.insert(
                        key,
                        json!({
                            "direct": types.first(),
                            "count": n_types,
                        }),
                    );
                }
            }
        }

        json!({
            "version": 1,
            "tier0": {
                "dir": "tier0",
                "broad_types": graph.broad_types(),
            },
            "tier1": tier1_models,
            "tier2": tier2_models,
            "tier2_min_types": self.config.tier2_min_types,
        })
    }

    /// Get the best available device.
    fn get_device() -> Device {
        #[cfg(feature = "cuda")]
        {
            if let Ok(device) = Device::new_cuda(0) {
                return device;
            }
        }

        #[cfg(feature = "metal")]
        {
            if let Ok(device) = Device::new_metal(0) {
                return device;
            }
        }

        Device::Cpu
    }
}

/// Result for a single tier model training.
#[derive(Debug, Clone)]
pub struct TierModelResult {
    pub name: String,
    pub classes: usize,
    pub accuracy: f64,
}

/// Report from tiered training.
#[derive(Debug, Clone, Default)]
pub struct TieredTrainingReport {
    pub tier0_accuracy: f64,
    pub tier0_classes: usize,
    pub tier1_results: Vec<TierModelResult>,
    pub tier1_skipped: Vec<String>,
    pub tier2_results: Vec<TierModelResult>,
}

impl std::fmt::Display for TieredTrainingReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Tiered Training Report:")?;
        writeln!(
            f,
            "  Tier 0: {:.2}% accuracy ({} classes)",
            self.tier0_accuracy * 100.0,
            self.tier0_classes
        )?;
        writeln!(
            f,
            "  Tier 1: {} models trained, {} skipped (single category)",
            self.tier1_results.len(),
            self.tier1_skipped.len()
        )?;
        for r in &self.tier1_results {
            writeln!(
                f,
                "    [{}]: {:.2}% ({} categories)",
                r.name,
                r.accuracy * 100.0,
                r.classes
            )?;
        }
        writeln!(f, "  Tier 2: {} models trained", self.tier2_results.len())?;
        for r in &self.tier2_results {
            writeln!(
                f,
                "    [{}]: {:.2}% ({} types)",
                r.name,
                r.accuracy * 100.0,
                r.classes
            )?;
        }
        Ok(())
    }
}