from beartype import beartype
import typing
import torch

@beartype
def _format_simple_assert(first: str, second: typing.Any, lt: bool) -> str:
    if lt:
        return f'(assert (<= {first} {second}))'
    return f'(assert (>= {first} {second}))'

@beartype
def _format_simple_declare(name: list | str, dtype: str) -> str:
    if isinstance(name, list):
        name = ' '.join(name)
    if dtype == 'Real':
        return f'(declare-const {name} {dtype})'
    return f'(declare-pwl {name} {dtype})'

@beartype
def _format_simple_proof(ap: list) -> str:
    # TODO: generalize to other activation
    def _format_node(value):
        assert isinstance(value, int) and value != 0, f'Invalid {value=}'
        if value > 0:
            return f'(>= N_{abs(value)} 0)'
        return f'(< N_{abs(value)} 0)'
    ret = [_format_node(_) for _ in ap]
    return f'(and {" ".join(ret)})'


@beartype
def _format_simple_proof_input(lower: torch.Tensor, upper: torch.Tensor) -> str:
    def _format_node(index, l, u) -> str:
        return f'(>= X_{index} {l}) (<= X_{index} {u})'
    ret = [_format_node(i, l, u) for i, (l, u) in enumerate(zip(lower.flatten(), upper.flatten()))]
    return f'(and {" ".join(ret)})'


@beartype
def create_aptp(
    proof: list, 
    input_lower: torch.Tensor, 
    input_upper: torch.Tensor, 
    cnf_cs: torch.Tensor, 
    cnf_rhs: torch.Tensor,
    input_split: bool,
) -> str:
    assert len(cnf_cs) == len(cnf_rhs)
    assert torch.all(input_lower <= input_upper)
    
    declare_lines, input_lines, output_lines, hidden_lines = [], [], [], []
    declare_lines.extend([_format_simple_declare(f'X_{i}', 'Real') for i in range(len(input_lower.flatten()))])
    declare_lines.extend([_format_simple_declare(f'Y_{i}', 'Real') for i in range(len(cnf_cs[0]))])
    if not input_split:
        # input
        for i, (l, u) in enumerate(zip(input_lower.flatten(), input_upper.flatten())):
            input_lines.append(_format_simple_assert(f'X_{i}', u, lt=True))
            input_lines.append(_format_simple_assert(f'X_{i}', l, lt=False))
        
        # hidden
        n_hiddens = max(abs(item) for path in proof for item in path)
        # TODO: generalize to other activation
        declare_lines.extend([_format_simple_declare(f'N_{i+1}', 'ReLU') for i in range(n_hiddens)]) 

        
    # output
    for (cs, rhs) in zip(cnf_cs, cnf_rhs):
        if (cs == 1).any():
            cs_ = cs
            rhs_ = rhs
            lt = True
        else:
            cs_ = cs * -1
            rhs_ = rhs * -1
            lt = False
            
        l = [f'Y_{i}' for i, v in enumerate(cs_) if v == 1]
        r = [f'Y_{i}' for i, v in enumerate(cs_) if v == -1]
        assert len(l)
        if rhs_ != 0.0 or not len(r):
            r.append(rhs_)
        assert len(l) == len(r) == 1, f'{cs=}, {rhs=}, {l=}, {r=}'
        output_lines.append(_format_simple_assert(l[0], r[0], lt=lt))
    
    # hidden
    if input_split:
        # raise NotImplementedError('Input split is not supported for APTP export')
        hidden_lines.extend(_format_simple_proof_input(l, u) for (l, u) in proof)
    else:
        hidden_lines.extend([_format_simple_proof(p) for p in proof])
    
    # combine
    aptp_str = '; APTP generated by NeuralSAT\n\n'
    aptp_str += '; Variables\n'
    aptp_str += '\n'.join(declare_lines) + '\n\n'
    if input_lines:
        aptp_str += '; Input constraints\n'
        aptp_str += '\n'.join(input_lines) + '\n\n'
    aptp_str += '; Output constraints\n'
    aptp_str += '\n'.join(output_lines) + '\n\n'
    aptp_str += '; Reasoning constraints\n'
    if hidden_lines:
        aptp_str += '(assert (or\n    ' + '\n    '.join(hidden_lines) + '\n))'
    return aptp_str

