import pandas as pd
import re
import sys
from allennlp.predictors import Predictor
import allennlp_models.tagging

# ✅ Load SRL model
predictor = Predictor.from_path(
    "https://storage.googleapis.com/allennlp-public-models/structured-prediction-srl-bert.2020.12.15.tar.gz"
)

# -----------------------
# Regex Habit Patterns
# -----------------------
HABIT_PATTERNS = re.compile(
    r"(i always|i usually|i tend to|i keep|i still|i never|"
    r"every day|every night|each morning|each night|all the time|every time|"
    r"without thinking|by default|automatically|out of habit|instinctively|"
    r"can't help|i end up|before bed|after work|when i wake up|as soon as i|get home)",
    re.IGNORECASE
)

# -----------------------
# Detect habits
# -----------------------
def detect_habit_regex(text):
    """Detect explicit habits using regex."""
    if not isinstance(text, str):
        return None
    matches = HABIT_PATTERNS.findall(text.lower())
    return "; ".join(set(matches)) if matches else None


def detect_habit_srl(text):
    """Detect habits using SRL (time anchors = habitual contexts)."""
    if not isinstance(text, str) or not text.strip():
        return None

    try:
        results = predictor.predict(sentence=text)
    except Exception:
        return None

    temporal_modifiers = []
    for verb in results["verbs"]:
        desc = verb["description"]
        spans = desc.split("]")
        for span in spans:
            if "[" not in span or ":" not in span:
                continue
            role, phrase = span.split(":", 1)
            role = role.strip("[ ")
            phrase = phrase.strip()
            if role == "ARGM-TMP":  # temporal modifier
                temporal_modifiers.append(phrase)

    return "; ".join(set(temporal_modifiers)) if temporal_modifiers else None


# -----------------------
# Main
# -----------------------
if __name__ == "__main__":
    infile = sys.argv[1]
    df = pd.read_csv(infile)

    # Apply habit detectors
    df["habit_regex"] = df["comment"].apply(detect_habit_regex)
    df["habit_srl"] = df["comment"].apply(detect_habit_srl)

    # Merge them: regex OR SRL evidence
    def merge_habits(row):
        habits = []
        if row["habit_regex"]:
            habits.append(row["habit_regex"])
        if row["habit_srl"]:
            habits.append(row["habit_srl"])
        return "; ".join(habits) if habits else None

    df["habit_detected"] = df.apply(merge_habits, axis=1)

    outfile = infile.replace(".csv", "_habits.csv")
    df.to_csv(outfile, index=False)
    print(f"✅ Saved habit detections to {outfile}")