#!/usr/bin/env python3
"""Abstract the LPC spec into boolean LTL for reactive synthesis.

The LPC spec uses real-valued variables (kias/kgs) and enumerated modes.
This script creates a boolean abstraction suitable for ltlsynt.

Speed regions (kias = kgs, so one variable):
  A: speed <= 20   (hover zone)
  B: 20 < speed <= 30
  C: 30 < speed <= 40
  D: 40 < speed <= 90
  E: 90 < speed <= 100
  F: speed > 100

Modes (lift_mode):
  tb:  thrust_borne
  stb: semi_thrust_borne
  swb: semi_wing_borne
  wb:  wing_borne

The environment controls speed (within adjacency constraints).
The controller chooses the mode.
"""

import json
from pathlib import Path


def build_lpc_spec():
    regions = ['a', 'b', 'c', 'd', 'e', 'f']
    modes = ['tb', 'stb', 'swb', 'wb']

    # Helper: exactly one of a list is true
    def exactly_one(vars):
        # At least one
        at_least = ' | '.join(vars)
        # At most one (pairwise exclusion)
        pairs = []
        for i in range(len(vars)):
            for j in range(i + 1, len(vars)):
                pairs.append(f'(!{vars[i]} | !{vars[j]})')
        at_most = ' & '.join(pairs)
        return f'({at_least}) & {at_most}'

    region_vars = [f'r_{r}' for r in regions]
    mode_vars = [f'm_{m}' for m in modes]

    # === ENVIRONMENT ASSUMPTIONS ===
    assumptions = []

    # Exactly one speed region (always)
    assumptions.append(f'G({exactly_one(region_vars)})')

    # Initial speed: region F (kias = 120 > 100)
    assumptions.append('r_f')

    # Speed adjacency: can only move to adjacent region per tick
    # (from the derivative constraint: |delta kias| <= 10 per tick)
    adjacency = {
        'r_a': ['r_a', 'r_b'],
        'r_b': ['r_a', 'r_b', 'r_c'],
        'r_c': ['r_b', 'r_c', 'r_d'],
        'r_d': ['r_c', 'r_d', 'r_e'],
        'r_e': ['r_d', 'r_e', 'r_f'],
        'r_f': ['r_e', 'r_f'],
    }
    for region, neighbors in adjacency.items():
        next_options = ' | '.join(f'X({n})' for n in neighbors)
        assumptions.append(f'G({region} -> ({next_options}))')

    # hover_control_mode <=> speed <= 20 (region A)
    # We absorb hover into r_a directly in the guarantees

    # === CONTROLLER GUARANTEES ===
    guarantees = []

    # Exactly one mode (always)
    guarantees.append(f'G({exactly_one(mode_vars)})')

    # Initial mode: wing_borne (since kias = 120 >= 90)
    # LPC_INIT_LIFT_MODE: lift_mode = wb <=> kias >= 90, and kias = 120
    guarantees.append('m_wb')

    # === MODE TRANSITIONS (from FRET requirements) ===

    # LPC_WB_STAY_ON_NEXT: While wb & kias > 90 -> next wb
    # kias > 90 means region E or F
    guarantees.append('G((m_wb & (r_e | r_f)) -> X(m_wb))')

    # LPC_WB_TO_SWB: Upon wb & kias <= 90 -> next swb
    # kias <= 90 means region A-D
    guarantees.append('G((m_wb & (r_a | r_b | r_c | r_d)) -> X(m_swb))')

    # LPC_SWB_STAY_ON_NEXT: While swb & kias <= 100 & kias > 30 -> next swb
    # 30 < kias <= 100 means region D or E
    guarantees.append('G((m_swb & (r_d | r_e)) -> X(m_swb))')

    # LPC_SWB_TO_WB: Upon swb & kias > 100 -> next wb
    # kias > 100 means region F
    guarantees.append('G((m_swb & r_f) -> X(m_wb))')

    # LPC_SWB_TO_STB: Upon swb & kias <= 30 -> next stb
    # kias <= 30 means region A or B
    guarantees.append('G((m_swb & (r_a | r_b)) -> X(m_stb))')

    # LPC_STB_STAY_ON_NEXT: While stb & kias <= 40 & (kgs > 20 | !hover)
    # kias <= 40 means region A-C; kgs > 20 | !hover means NOT region A
    # So: stb & (region B or C) -> next stb
    guarantees.append('G((m_stb & (r_b | r_c)) -> X(m_stb))')

    # LPC_STB_TO_SWB: Upon stb & kias > 40 -> next swb
    # kias > 40 means region D-F
    guarantees.append('G((m_stb & (r_d | r_e | r_f)) -> X(m_swb))')

    # LPC_STB_TO_TB: Upon stb & hover & kgs <= 20 -> next tb
    # hover & kgs <= 20 means region A
    guarantees.append('G((m_stb & r_a) -> X(m_tb))')

    # LPC_TB_STAY_ON_NEXT: While tb & kgs <= 20 & hover -> next tb
    # region A
    guarantees.append('G((m_tb & r_a) -> X(m_tb))')

    # LPC_TB_TO_STB: Upon tb & (!hover | kgs > 20) -> next stb
    # NOT region A
    guarantees.append('G((m_tb & (r_b | r_c | r_d | r_e | r_f)) -> X(m_tb | m_stb))')
    # Wait, the original says -> next stb. Let me fix:
    # Actually: Upon tb & (!hover | kgs > 20) -> next stb
    # !hover | kgs > 20 means NOT in region A
    guarantees[-1] = 'G((m_tb & !r_a) -> X(m_stb))'

    # Note: we skip LPC_REACH_HOVER_06 and _12 (bounded temporal F[0,n])
    # ltlsynt doesn't support bounded-F natively

    # === BUILD FORMULA ===
    # assume-guarantee: assumptions -> guarantees
    assume_str = ' & '.join(f'({a})' for a in assumptions)
    guarantee_str = ' & '.join(f'({g})' for g in guarantees)
    formula = f'({assume_str}) -> ({guarantee_str})'

    inputs = region_vars
    outputs = mode_vars

    return {
        '_comment': 'Boolean abstraction of LPC lift-mode controller. Assumes -> Guarantees form.',
        'spec_name': 'LPC_mini',
        'source_file': 'specs/LPC_mini_reqts_and_vars.json',
        'inputs': inputs,
        'outputs': outputs,
        'assumptions': assumptions,
        'guarantees': guarantees,
        'conjoined_ltl': formula,
        'abstraction_notes': {
            'speed_regions': {
                'r_a': 'speed <= 20 (hover zone)',
                'r_b': '20 < speed <= 30',
                'r_c': '30 < speed <= 40',
                'r_d': '40 < speed <= 90',
                'r_e': '90 < speed <= 100',
                'r_f': 'speed > 100',
            },
            'modes': {
                'm_tb': 'thrust_borne',
                'm_stb': 'semi_thrust_borne',
                'm_swb': 'semi_wing_borne',
                'm_wb': 'wing_borne',
            },
            'skipped_requirements': [
                'LPC_REACH_HOVER_06 (bounded temporal F[0,6])',
                'LPC_REACH_HOVER_12 (bounded temporal F[0,12])',
                'LPC_KIAS_DERIVATIVE (absorbed into adjacency constraint)',
                'LPC_KIAS_KGS (absorbed - single speed variable)',
                'LPC_INIT_KIAS (absorbed into initial region)',
                'LPC_KIAS_0 (trivially satisfied)',
                'LPC_INIT_HOVER_MODE (absorbed - hover = region A)',
            ],
        },
    }


if __name__ == '__main__':
    config = build_lpc_spec()
    out = Path('specs/lpc_synthesis_config.json')
    out.write_text(json.dumps(config, indent=2) + '\n')

    print(f'LPC Boolean Abstraction')
    print(f'Inputs  ({len(config["inputs"])}): {", ".join(config["inputs"])}')
    print(f'Outputs ({len(config["outputs"])}): {", ".join(config["outputs"])}')
    print(f'\nAssumptions ({len(config["assumptions"])}):')
    for a in config['assumptions']:
        print(f'  {a}')
    print(f'\nGuarantees ({len(config["guarantees"])}):')
    for g in config['guarantees']:
        print(f'  {g}')
    print(f'\nSkipped: {len(config["abstraction_notes"]["skipped_requirements"])} requirements')
    print(f'Config written to: {out}')
