#!/usr/bin/env python3
"""Convert FRET JSON export to ltlsynt synthesis config.

Automatically detects:
  - Mode variables: any variable matching `control_* = q_*` in the LTL.
    Each control group (e.g. control_mode) gets its own set of boolean
    outputs (e.g. in_mode_shutdown, in_mode_heatup) with a mutual
    exclusion constraint.
  - Environment inputs: variables that appear only in conditions
    (pre_condition / regular_condition) and are not mode variables.
  - System outputs: the boolean mode propositions derived from control_*.

Naming convention in FRET:
  - Mode variables MUST start with "control_" (e.g. control_mode, control_pump)
  - Mode values MUST start with "q_" (e.g. q_shutdown, q_heatup)
  - All other variables are classified by their FRET role.

Usage: python3 scripts/fret_to_synth.py <fret_export.json> [output.json]
"""

import json
import re
import sys
from collections import defaultdict
from pathlib import Path


def discover_modes(requirements):
    """Scan all LTL formulas for (control_X = q_Y) patterns.

    Returns a dict: { "control_X": ["q_Y", "q_Z", ...] }
    Each key is a mode group, values are the discovered mode values.
    """
    # Match patterns like (control_mode = q_shutdown), (control_pump = q_off)
    pattern = re.compile(r'\((\s*control_\w+)\s*=\s*(q_\w+)\s*\)')

    mode_groups = defaultdict(set)
    for req in requirements:
        sem = req["semantics"]
        for field in ("ftInfAUExpanded", "ft", "ftExpanded"):
            ltl = sem.get(field, "")
            if ltl:
                for match in pattern.finditer(ltl):
                    ctrl_var = match.group(1).strip()
                    mode_val = match.group(2).strip()
                    mode_groups[ctrl_var].add(mode_val)

    # Sort mode values for deterministic output
    return {k: sorted(v) for k, v in sorted(mode_groups.items())}


def build_mode_map(mode_groups):
    """Build the string replacement map and output variable list.

    For a single control group like control_mode with values
    [q_heatup, q_operation, q_scram, q_shutdown], produces:
      - replacements: {"(control_mode = q_heatup)": "in_mode_heatup", ...}
      - outputs: ["in_mode_heatup", "in_mode_operation", ...]

    For multiple groups, the boolean name includes the group suffix:
      control_mode   -> in_mode_X
      control_pump   -> in_pump_X
    """
    replacements = {}
    outputs_by_group = {}

    for ctrl_var, mode_vals in mode_groups.items():
        # Derive group suffix: control_mode -> mode, control_pump -> pump
        suffix = ctrl_var.replace("control_", "", 1)
        group_outputs = []

        for qval in mode_vals:
            # q_shutdown -> shutdown
            val_name = qval.replace("q_", "", 1)
            bool_name = f"in_{suffix}_{val_name}"
            replacements[f"({ctrl_var} = {qval})"] = bool_name
            group_outputs.append(bool_name)

        outputs_by_group[ctrl_var] = group_outputs

    return replacements, outputs_by_group


def classify_variables(requirements, mode_groups):
    """Classify non-mode variables as environment inputs.

    A variable is an environment input if:
      - It is NOT a control_* or q_* variable (those are handled by mode encoding)
      - It appears in any requirement's semantics.variables

    All mode propositions are system outputs (handled separately).
    """
    # Collect all known mode-related variable names to exclude
    mode_var_names = set()
    for ctrl_var, mode_vals in mode_groups.items():
        mode_var_names.add(ctrl_var)
        mode_var_names.update(mode_vals)

    # Collect all variables referenced across requirements
    all_vars = set()
    for req in requirements:
        sem = req["semantics"]
        for v in sem.get("variables", []):
            all_vars.add(v)

    # Everything that's not a mode variable is an environment input
    env_inputs = sorted(all_vars - mode_var_names)
    return env_inputs


def encode_ltl(formula, replacements):
    """Replace all (control_X = q_Y) with boolean mode propositions."""
    result = formula
    for pattern, replacement in replacements.items():
        result = result.replace(pattern, replacement)

    # Sanity check: no control_ or q_ references should remain
    leftover = re.findall(r'control_\w+|q_\w+', result)
    if leftover:
        raise ValueError(
            f"Unhandled mode references in LTL: {leftover}\n"
            f"  Formula: {result}\n"
            f"  Known replacements: {list(replacements.keys())}"
        )

    return result


def mutual_exclusion_ltl(group_outputs):
    """Generate LTL for exactly-one-active constraint over a mode group.

    Produces: G( (a & !b & !c & !d) | (!a & b & !c & !d) | ... )
    """
    clauses = []
    for i, m in enumerate(group_outputs):
        parts = []
        for j, n in enumerate(group_outputs):
            if i == j:
                parts.append(n)
            else:
                parts.append(f"(! {n})")
        clauses.append(f"({' & '.join(parts)})")
    exactly_one = " | ".join(clauses)
    return f"G ({exactly_one})"


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <fret_export.json> [output.json]", file=sys.stderr)
        sys.exit(1)

    input_path = Path(sys.argv[1])
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("specs/synthesis_config_v3.json")

    with open(input_path) as f:
        fret_data = json.load(f)

    requirements = fret_data["requirements"]
    project = requirements[0]["project"]
    component = requirements[0]["semantics"]["component_name"]

    # --- Auto-discover everything ---
    mode_groups = discover_modes(requirements)
    replacements, outputs_by_group = build_mode_map(mode_groups)
    env_inputs = classify_variables(requirements, mode_groups)
    all_outputs = [out for group in outputs_by_group.values() for out in group]

    print(f"Discovered {len(mode_groups)} mode group(s):")
    for ctrl_var, modes in mode_groups.items():
        print(f"  {ctrl_var}: {modes}")
        print(f"    -> booleans: {outputs_by_group[ctrl_var]}")
    print(f"Environment inputs: {env_inputs}")
    print(f"System outputs: {all_outputs}")
    print()

    # --- Encode requirements ---
    encoded_reqs = []
    for req in requirements:
        reqid = req["reqid"]
        fulltext = req["fulltext"]
        sem = req["semantics"]

        # Prefer ftInfAUExpanded (cleanest infinite-horizon LTL)
        raw_ltl = sem.get("ftInfAUExpanded") or sem.get("ft")
        if not raw_ltl:
            print(f"WARNING: No LTL for {reqid}, skipping", file=sys.stderr)
            continue

        encoded = encode_ltl(raw_ltl, replacements)

        encoded_reqs.append({
            "req_id": reqid,
            "fulltext": fulltext,
            "project": project,
            "component": component,
            "ltl": encoded,
            "ltl_original": raw_ltl,
            "condition_type": sem.get("condition", "unknown"),
        })

    # --- Build structural constraints (one mutex per mode group) ---
    structural = []
    for ctrl_var, group_outs in outputs_by_group.items():
        mutex = mutual_exclusion_ltl(group_outs)
        structural.append({
            "name": f"mutex_{ctrl_var}",
            "description": f"Exactly one {ctrl_var} value active at all times",
            "ltl": mutex,
        })

    # --- Conjoin all formulas ---
    all_formulas = [r["ltl"] for r in encoded_reqs]
    all_formulas += [s["ltl"] for s in structural]
    conjoined = " & ".join(f"({f})" for f in all_formulas)

    # --- Write config ---
    config = {
        "_comment": (
            f"Generated from {input_path.name} by fret_to_synth.py. "
            f"Boolean mode encoding with auto-discovered mode groups. "
            f"Mutual exclusion constraints are synthesis artifacts — they enforce "
            f"the single-valued semantics of control_* variables after decomposition "
            f"into independent booleans. They are NOT needed in FRET itself, where "
            f"control_mode = q_X is inherently single-valued."
        ),
        "spec_name": f"{project}_{component}",
        "source_file": str(input_path),
        "mode_groups": {
            ctrl_var: {
                "values": modes,
                "booleans": outputs_by_group[ctrl_var],
            }
            for ctrl_var, modes in mode_groups.items()
        },
        "inputs": env_inputs,
        "outputs": all_outputs,
        "requirements": encoded_reqs,
        "structural_constraints": structural,
        "conjoined_ltl": conjoined,
    }

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(config, f, indent=2)

    print(f"Synthesis config written to: {output_path}")
    print(f"  Requirements: {len(encoded_reqs)}")
    print(f"  Structural constraints: {len(structural)}")
    print(f"  Conjoined formula length: {len(conjoined)} chars")

    print("\nEncoded requirements:")
    for r in encoded_reqs:
        print(f"  {r['req_id']} [{r['condition_type']}]:")
        print(f"    {r['ltl'][:120]}{'...' if len(r['ltl']) > 120 else ''}")


if __name__ == "__main__":
    main()