#
#   Imandra Inc.
#
#   speclogician/state/diff.py
#

from enum import StrEnum
from typing import TypeAlias, Callable, TypeVar, Generic
from pydantic import BaseModel, Field

from ..modeling.domain import BaseStatus
from ..modeling.complement import ScenarioComplementDiff
from ..modeling.reachability import ScenarioReachabilityGraphDiff
from ..modeling.trace_reachability import TraceReachabilityDiff


class ComparisonOutcome(StrEnum):
    """ When we compare the before/after values, we'll use this as the comparison outcome """
    UNKNOWN = "unknown"
    NO_CHANGE = "no_change"
    NO_CHANGE_GOOD = "no_change_good"
    NO_CHANGE_BAD = "no_change_bad"
    IMPROVED = "improved"
    DECLINED = "declined"


T = TypeVar("T")

ComparisonFn: TypeAlias = Callable[[T, T], ComparisonOutcome]

class ValueDiff(BaseModel, Generic[T]):
    """ Represents a change (before → after) with an explanatory label
    and an interpretation of whether increases/decreases are “good”. """

    label : str 
    before: T
    after: T
    comp_func : ComparisonFn[T] | None = Field(default=None, exclude=True)
    hint : None | dict[ComparisonOutcome, str] = None

    def is_different(self) -> bool:
        if self.comp_func:
            return self.comp_func(self.before, self.after) in (ComparisonOutcome.IMPROVED, ComparisonOutcome.DECLINED)
        return self.before != self.after


# comparison functions 

def base_status_comp (before : BaseStatus, after : BaseStatus) -> ComparisonOutcome:
    """ """
    # TODO there's prob a nice way to do this without explicit write out...
    match (before, after):
        case (BaseStatus.UNKNOWN, BaseStatus.UNKNOWN):
            return ComparisonOutcome.NO_CHANGE_BAD
        case (BaseStatus.UNKNOWN, _):
            return ComparisonOutcome.IMPROVED
        case (BaseStatus.INVALID_IML, BaseStatus.VALID):
            return ComparisonOutcome.IMPROVED
        case (BaseStatus.INVALID_IML, BaseStatus.INVALID_IML):
            return ComparisonOutcome.NO_CHANGE_BAD
        case (BaseStatus.VALID, BaseStatus.VALID):
            return ComparisonOutcome.NO_CHANGE_GOOD
        case (BaseStatus.VALID, _):
            return ComparisonOutcome.DECLINED
        case (_, _):
            return ComparisonOutcome.UNKNOWN

def bool_true_is_good (before : bool, after : bool) -> ComparisonOutcome:
    """ Compare two boolean values where `true` is good """
    match (before, after):
        case (False, False):
            return ComparisonOutcome.NO_CHANGE_BAD
        case (False, True):
            return ComparisonOutcome.IMPROVED
        case (True, False):
            return ComparisonOutcome.DECLINED
        case (True, True):
            return ComparisonOutcome.NO_CHANGE_GOOD

def bool_true_is_bad(before: bool, after: bool) -> ComparisonOutcome:
    match (before, after):
        case (False, False):
            return ComparisonOutcome.NO_CHANGE_GOOD
        case (False, True):
            return ComparisonOutcome.DECLINED
        case (True, False):
            return ComparisonOutcome.IMPROVED
        case (True, True):
            return ComparisonOutcome.NO_CHANGE_BAD

def numeric_increase_good(before: int | float, after: int | float) -> ComparisonOutcome:
    """
    Increase is good (e.g. matched counts).
    """
    if before == after:
        return ComparisonOutcome.NO_CHANGE_GOOD if before > 0 else ComparisonOutcome.NO_CHANGE
    elif before < after:
        return ComparisonOutcome.IMPROVED
    else:
        return ComparisonOutcome.DECLINED


def numeric_increase_bad(before: int | float, after: int | float) -> ComparisonOutcome:
    """
    Increase is bad (e.g. errors, inconsistencies).
    """
    if before == after:
        return ComparisonOutcome.NO_CHANGE_BAD if before > 0 else ComparisonOutcome.NO_CHANGE_GOOD
    elif before < after:
        return ComparisonOutcome.DECLINED
    else:
        return ComparisonOutcome.IMPROVED

def complement_regions_decrease_good(before: int | None, after: int | None) -> ComparisonOutcome:
    """
    Complement represents uncovered space.
    Smaller complement => improved coverage.
    """
    if before is None or after is None:
        return ComparisonOutcome.UNKNOWN
    if after < before:
        return ComparisonOutcome.IMPROVED
    if after > before:
        return ComparisonOutcome.DECLINED
    return ComparisonOutcome.NO_CHANGE


def complement_present_comp(before: bool, after: bool) -> ComparisonOutcome:
    """
    Presence itself isn't strictly good/bad; treat as UNKNOWN unless equal.
    (We separately judge region count where smaller is better.)
    """
    if before == after:
        return ComparisonOutcome.NO_CHANGE
    return ComparisonOutcome.UNKNOWN


class StateDiff(BaseModel):
    """ Statistical summary of changes between two states """

    # Domain model-related value diffs
    base_status : ValueDiff[BaseStatus]         # Base status
    base_has_state : ValueDiff[bool]            # Does the base have state defined?
    base_has_action : ValueDiff[bool]           # Does the base have action defined?

    # State predicates
    num_state_preds_total : ValueDiff[int]      # 
    num_state_preds_matched : ValueDiff[int]    # 
    num_state_preds_valid_logic : ValueDiff[int] #
    num_state_preds_errored : ValueDiff[int]    #

    # state x action predicates
    num_action_preds_total : ValueDiff[int]     #
    num_action_preds_matched : ValueDiff[int]   #
    num_action_preds_valid_logic : ValueDiff[int] #
    num_action_preds_errored : ValueDiff[int]   #

    # Predicates
    num_preds_total : ValueDiff[int]            # Total number of predicates
    num_preds_matched : ValueDiff[int]          # Predicates matched
    num_preds_valid_logic : ValueDiff[int]      # Number of predicates with valid logic
    num_preds_errored : ValueDiff[int]          # Number of predicates with errors
    
    # Transitions
    num_trans_total : ValueDiff[int]            # Total number 
    num_trans_matched : ValueDiff[int]          # Number matched
    num_trans_valid_logic : ValueDiff[int]      # Number with valid logic
    num_trans_errored : ValueDiff[int]          # Number of transitions with errors

    # Scenarios
    num_sc_total : ValueDiff[int]               # Total number of scenarios
    num_sc_missing : ValueDiff[int]             # Missing predicates/transitions
    num_sc_matched : ValueDiff[int]             # Scenarios matched vs artifacts
    num_sc_inconsistent : ValueDiff[int]        # Scenarios inconsistent

    # Conflicts (spec-level)
    num_sc_conflicted : ValueDiff[int]          # Total conflicts (all kinds)
    num_sc_overlap : ValueDiff[int]             # Overlap conflicts
    num_sc_consumed : ValueDiff[int]            # Consumed conflicts

    # Test traces
    num_test_traces_total : ValueDiff[int]      # Total number of test traces
    num_test_traces_logic_good : ValueDiff[int] # Valid logic
    num_test_traces_matched : ValueDiff[int]    # Matched traces

    # Log traces
    num_log_traces_total : ValueDiff[int]       # Total number of log traces
    num_log_traces_logic_good : ValueDiff[int]  # Number of log traces with valid logic
    num_log_traces_matched : ValueDiff[int]     # Number of matched traces

    # Scenario complement
    scenario_comp_present: ValueDiff[bool]
    num_comp_regions_total: ValueDiff[int | None]
    scenario_comp_diff: ScenarioComplementDiff | None = None

    # Reachability graph
    scenario_reach_present: ValueDiff[bool]
    reach_num_nodes: ValueDiff[int]
    reach_num_edges: ValueDiff[int]
    reach_num_initials: ValueDiff[int]
    reach_num_reachable: ValueDiff[int]
    reach_num_unreachable: ValueDiff[int]

    scenario_reach_diff: ScenarioReachabilityGraphDiff | None = None

    # ----------------------------
    # Trace reachability analysis
    # ----------------------------

    trace_reach_present: ValueDiff[bool]

    tr_total_traces: ValueDiff[int]
    tr_total_steps: ValueDiff[int]
    tr_total_hops: ValueDiff[int]

    tr_supported_hops: ValueDiff[int]
    tr_unsupported_hops: ValueDiff[int]

    tr_scenarios_hit: ValueDiff[int]
    tr_scenarios_unhit: ValueDiff[int]

    trace_reach_diff: TraceReachabilityDiff | None = None