"""
Baseline: Dictionary + TF-IDF retrieval for Akkadian-to-English translation.
This is a Kaggle submission script (no GPU needed).
"""
import os
import re
import numpy as np
import pandas as pd
from collections import Counter
from sklearn.feature_extraction.text import TfidfVectorizer

# === Paths (Kaggle environment) ===
# Debug: find the actual data location
print("Searching for data files...")
INPUT_DIR = None
for candidate in [
    "/kaggle/input/deep-past-initiative-machine-translation",
    "/kaggle/input/competitions/deep-past-initiative-machine-translation",
]:
    if os.path.exists(candidate):
        INPUT_DIR = candidate
        break

if INPUT_DIR is None:
    # Walk /kaggle/input to find train.csv
    for root, dirs, files in os.walk("/kaggle/input"):
        print(f"  {root}: {files[:10]}")
        if "train.csv" in files:
            INPUT_DIR = root
            break

if INPUT_DIR is None:
    raise FileNotFoundError("Cannot find competition data! Listing /kaggle/input:")

print(f"Using INPUT_DIR: {INPUT_DIR}")
for f in sorted(os.listdir(INPUT_DIR)):
    fpath = os.path.join(INPUT_DIR, f)
    size = os.path.getsize(fpath) if os.path.isfile(fpath) else 0
    print(f"  {f} ({size} bytes)")

train_df = pd.read_csv(f"{INPUT_DIR}/train.csv")
test_df = pd.read_csv(f"{INPUT_DIR}/test.csv")
sentences_df = pd.read_csv(f"{INPUT_DIR}/Sentences_Oare_FirstWord_LinNum.csv")
published_df = pd.read_csv(f"{INPUT_DIR}/published_texts.csv")

print(f"Train: {len(train_df)}, Test: {len(test_df)}, Sentences: {len(sentences_df)}")

# === Build translation memory from train data ===
memory = []
for _, row in train_df.iterrows():
    src = str(row['transliteration']).strip()
    tgt = str(row['translation']).strip()
    if src and tgt:
        memory.append({'src': src, 'tgt': tgt})

# === Add sentence-level data from Sentences file ===
# Match sentences to published_texts transliterations
pub_translit = {}
for _, row in published_df.iterrows():
    uid = str(row['oare_id'])
    t = str(row.get('transliteration', ''))
    if t and t != 'nan':
        pub_translit[uid] = t

# Group sentences by text
sent_groups = sentences_df.groupby('text_uuid')
for text_uuid, group in sent_groups:
    group_sorted = group.sort_values('sentence_obj_in_text')
    for _, srow in group_sorted.iterrows():
        trans = str(srow.get('translation', ''))
        spelling = str(srow.get('first_word_spelling', ''))
        if trans and trans != 'nan' and len(trans) > 3:
            # Use first_word_spelling as a minimal source hint
            if spelling and spelling != 'nan':
                memory.append({'src': spelling, 'tgt': trans})

print(f"Translation memory: {len(memory)} entries")

# === Normalize source text ===
def norm_src(text):
    text = str(text).lower().strip()
    text = re.sub(r'[₀-₉0-9]', '', text)
    text = re.sub(r'\{[^}]*\}', ' ', text)
    text = re.sub(r'[<>\[\]]', ' ', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

mem_srcs_raw = [m['src'] for m in memory]
mem_srcs_norm = [norm_src(s) for s in mem_srcs_raw]
mem_tgts = [m['tgt'] for m in memory]

# === Build TF-IDF index on memory sources ===
char_vec = TfidfVectorizer(analyzer='char_wb', ngram_range=(3, 5), max_features=50000)
word_vec = TfidfVectorizer(analyzer='word', ngram_range=(1, 2), max_features=30000)

Xc = char_vec.fit_transform(mem_srcs_norm)
Xw = word_vec.fit_transform(mem_srcs_norm)

mem_lens = np.array([len(s.split()) for s in mem_srcs_norm])

# === Build word dictionary from train pairs ===
src_cnt, tgt_cnt, pair_cnt = Counter(), Counter(), Counter()
STOP_EN = {"the", "a", "an", "and", "or", "but", "of", "at", "by", "from", "with", "to", "in", "is", "for"}

for m in memory[:len(train_df)]:  # Only use train pairs for dictionary
    s_tokens = set(norm_src(m['src']).split())
    t_tokens = set(w for w in m['tgt'].lower().split() if w not in STOP_EN)
    for sw in s_tokens:
        src_cnt[sw] += 1
    for tw in t_tokens:
        tgt_cnt[tw] += 1
    for sw in s_tokens:
        for tw in t_tokens:
            pair_cnt[(sw, tw)] += 1

word_dict = {}
for (sw, tw), c in pair_cnt.items():
    score = 2 * c / (src_cnt[sw] + tgt_cnt[tw])
    if sw not in word_dict or score > word_dict[sw][1]:
        word_dict[sw] = (tw, score)

print(f"Word dictionary: {len(word_dict)} entries")

# === Translation function ===
def translate(query):
    q_norm = norm_src(query)
    if not q_norm:
        return "broken text"

    # TF-IDF retrieval
    qc = char_vec.transform([q_norm])
    qw = word_vec.transform([q_norm])

    sc = (qc @ Xc.T).toarray()[0]
    sw = (qw @ Xw.T).toarray()[0]

    combined = 0.7 * sc + 0.3 * sw

    # Length penalty
    q_len = max(1, len(q_norm.split()))
    ratios = mem_lens.astype(np.float32) / q_len
    lp = np.exp(-np.abs(np.log(ratios + 1e-5)) * 1.5)
    combined *= lp

    best_idx = np.argmax(combined)
    best_score = combined[best_idx]

    if best_score > 0.15:
        return mem_tgts[best_idx]
    else:
        # Fallback: word-by-word dictionary
        words = q_norm.split()
        result = []
        for w in words:
            if w in word_dict and word_dict[w][1] > 0.1:
                result.append(word_dict[w][0])
            else:
                result.append(w)
        return " ".join(result)

# === Generate predictions ===
print("Translating test set...")
predictions = []
for _, row in test_df.iterrows():
    pred = translate(str(row['transliteration']))
    if not pred or len(pred) == 0:
        pred = "broken text"
    predictions.append(pred)

# === Save submission ===
submission = pd.DataFrame({
    'id': test_df['id'],
    'translation': predictions
})
submission.to_csv('submission.csv', index=False)
print(f"Submission saved: {len(submission)} rows")
print(submission.head())