#!/usr/bin/env python3
import sqlite3
from transformers import MarianMTModel, MarianTokenizer, pipeline

print("Loading translation model...")
model_name = "Helsinki-NLP/opus-mt-nso-en"
tok = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

print("Loading sentiment model...")
classifier = pipeline("text-classification",
                      model="distilbert-base-uncased-finetuned-sst-2-english",
                      device=-1)
print("Models ready.\n")

def translate_batch(texts):
    inputs = tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=128)
    outputs = model.generate(**inputs, max_length=128)
    return [tok.decode(o, skip_special_tokens=True) for o in outputs]

conn = sqlite3.connect("/home/sediba/sepedi_datasets/sepedi_data.db")
done = set(r[0] for r in conn.execute("SELECT text FROM annotated_data").fetchall())

rows = conn.execute("""
    SELECT id, text, source FROM raw_texts
    WHERE length(text) BETWEEN 20 AND 250
    AND text NOT LIKE '%<%'
    AND text NOT LIKE '%>%'
    AND source IN ('jw.org', 'SADILAR-NCHLT-Annotated-Corpus')
    ORDER BY source DESC, RANDOM()
    LIMIT 2000
""").fetchall()

candidates = [(i, t, s) for i, t, s in rows if t not in done]
print(f"Candidates: {len(candidates)}")

inserted, errors = 0, 0
BATCH = 16

for i in range(0, len(candidates), BATCH):
    batch = candidates[i:i+BATCH]
    texts = [r[1] for r in batch]
    try:
        eng = translate_batch(texts)
        sentiments = classifier(eng, truncation=True, max_length=128)
        for (_, orig, src), sent in zip(batch, sentiments):
            label = sent["label"].lower()
            conn.execute("""
                INSERT OR IGNORE INTO annotated_data
                (text, label, label_name, annotator, confidence, date_annotated)
                VALUES (?, ?, ?, ?, ?, datetime('now'))
            """, (orig, 1 if label=="positive" else 0, label,
                  "auto:opus-mt+distilbert", round(sent["score"], 4)))
            inserted += 1
        conn.commit()
        if i % (BATCH*5) == 0:
            print(f"  {inserted} labeled...")
    except Exception as e:
        errors += 1
        print(f"  batch error: {e}")

total = conn.execute("SELECT COUNT(*) FROM annotated_data").fetchone()[0]
dist  = conn.execute("SELECT label_name, COUNT(*) FROM annotated_data GROUP BY label_name").fetchall()
print(f"\nDone. Inserted={inserted} Errors={errors} Total={total}")
for l, c in dist: print(f"  {l}: {c}")
conn.close()