// data.sy -- Training Data Pipeline
// Loads hormone-tagged memories from JSONL, tokenizes, batches.
//
// Marjorie McCubbins, Aislinn McCubbins, April 2026
// Patent Pending: 63/939,190 & 63/962,385 & 63/988,485

// === TRAINING SAMPLE ===

struct Sample {
    input_ids: Tensor,     // [seq_len] token IDs
    target_ids: Tensor,    // [seq_len] shifted by 1
    entity_id: Int,        // Who is this interaction with
    hormones: Vec<Float>,  // 5D hormone state at encoding time
    trust: Float,          // Trust level at encoding time
}

// === BATCH ===

struct Batch {
    input_ids: Tensor,     // [batch_size, seq_len]
    target_ids: Tensor,    // [batch_size, seq_len]
    entity_ids: Tensor,    // [batch_size]
    hormones: Tensor,      // [batch_size, 5]
    trust: Tensor,         // [batch_size]
}

// === DATA LOADER ===

struct DataLoader {
    samples: Vec<Sample>,
    batch_size: Int,
    current: Int,
    max_seq_len: Int,
    shuffle_seed: Int,
}

impl DataLoader {
    fn from_jsonl(path: String, max_seq_len: Int, batch_size: Int) -> DataLoader {
        let raw = read_file(path)
        let lines = str_split_lines(raw)
        let samples = []

        for i in 0..lines.len() {
            let line = lines[i]
            if str_len(line) < 5 { continue }

            // Parse JSONL line for prompt/response
            let prompt = json_get(line, "prompt")
            let response = json_get(line, "response")
            if str_len(prompt) < 1 { continue }

            // Tokenize
            let text = prompt + " " + response
            let tokens = tokenize(text)

            // Truncate to max_seq_len
            let seq_len = t_len(tokens)
            if seq_len > max_seq_len {
                tokens = t_slice(tokens, 0, max_seq_len)
                seq_len = max_seq_len
            }
            if seq_len < 4 { continue }

            // Input = tokens[0..n-1], target = tokens[1..n]
            let input = t_slice(tokens, 0, seq_len - 1)
            let target = t_slice(tokens, 1, seq_len)

            // Extract hormone state if present (default to tonic)
            let hormones = [
                json_get_float(line, "oxytocin", tonic_drive("oxytocin")),
                json_get_float(line, "cortisol", tonic_drive("cortisol")),
                json_get_float(line, "serotonin", tonic_drive("serotonin")),
                json_get_float(line, "dopamine", tonic_drive("dopamine")),
                json_get_float(line, "norepinephrine", tonic_drive("norepinephrine")),
            ]
            let trust_val = json_get_float(line, "trust", tonic_drive("oxytocin"))
            let eid = json_get_int(line, "entity_id", 0)

            samples.push(Sample {
                input_ids: input,
                target_ids: target,
                entity_id: eid,
                hormones: hormones,
                trust: trust_val,
            })
        }

        println("Loaded " + samples.len() + " samples from " + path)

        DataLoader {
            samples: samples,
            batch_size: batch_size,
            current: 0,
            max_seq_len: max_seq_len,
            shuffle_seed: 42,
        }
    }

    fn len(&self) -> Int {
        self.samples.len()
    }

    fn num_batches(&self) -> Int {
        self.samples.len() / self.batch_size
    }

    fn next_batch(mut self) -> Batch {
        let start = self.current
        let end = start + self.batch_size
        if end > self.samples.len() {
            self.current = 0
            start = 0
            end = self.batch_size
        }
        self.current = end

        // Collate samples into batch tensors
        let inputs = []
        let targets = []
        let entities = []
        let horm_list = []
        let trust_list = []

        for i in start..end {
            let s = self.samples[i]
            inputs.push(s.input_ids)
            targets.push(s.target_ids)
            entities.push(s.entity_id)
            horm_list.push(s.hormones)
            trust_list.push(s.trust)
        }

        // Stack into batch tensors
        Batch {
            input_ids: t_stack(inputs),
            target_ids: t_stack(targets),
            entity_ids: tensor_from_ints(entities),
            hormones: t_stack_floats(horm_list),
            trust: tensor_from_floats(trust_list),
        }
    }

    fn reset(mut self) {
        self.current = 0
    }
}
