// Copyright 2023-2024 The Regents of the University of California
// Copyright 2024 Cornell University
// released under BSD 3-Clause License
// author: Kevin Laeufer <laeufer@cornell.edu>

use patronus::btor2;
use patronus::expr::*;
use patronus::system::transform::{replace_anonymous_inputs_with_zero, simplify_expressions};

const COUNT_2: &str = r#"
1 sort bitvec 3
2 zero 1
3 state 1
4 init 1 3 2
5 one 1
6 add 1 3 5
7 next 1 3 6
8 ones 1
9 sort bitvec 1
10 eq 9 3 8
11 bad 10
"#;

#[test]
fn parse_count2() {
    let mut ctx = Context::default();
    let sys = btor2::parse_str(&mut ctx, COUNT_2, Some("count2")).unwrap();
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn serialize_count2() {
    let mut ctx = Context::default();
    let sys = btor2::parse_str(&mut ctx, COUNT_2, Some("count2")).unwrap();
    insta::assert_snapshot!(skip_first_line(&btor2::serialize_to_str(&ctx, &sys)));
}

const PASS_THROUGH: &str = r#"
1 sort bitvec 1
2 input 1 i
3 output 2 o
"#;

#[test]
fn serialize_pass_through() {
    let mut ctx = Context::default();
    let sys = btor2::parse_str(&mut ctx, PASS_THROUGH, Some("pass")).unwrap();
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn parse_quiz1() {
    let (ctx, sys) = btor2::parse_file("../inputs/chiseltest/Quiz1.btor").unwrap();
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn serialize_quiz1() {
    let (ctx, sys) = btor2::parse_file("../inputs/chiseltest/Quiz1.btor").unwrap();
    insta::assert_snapshot!(skip_first_line(&btor2::serialize_to_str(&ctx, &sys)));
}

#[test]
fn parse_instrumented_decoder() {
    let (ctx, sys) =
        btor2::parse_file("../inputs/repair/decoder_3_to_8.instrumented.btor").unwrap();
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn parse_sdram() {
    let (mut ctx, mut sys) =
        btor2::parse_file("../inputs/repair/sdram_controller.original.btor").unwrap();
    replace_anonymous_inputs_with_zero(&mut ctx, &mut sys);
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn parse_sdram_and_remove_anonymous_inputs() {
    let (mut ctx, mut sys) =
        btor2::parse_file("../inputs/repair/sdram_controller.original.btor").unwrap();
    replace_anonymous_inputs_with_zero(&mut ctx, &mut sys);
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn parse_sdram_and_simplify_expressions() {
    let (mut ctx, mut sys) =
        btor2::parse_file("../inputs/repair/sdram_controller.original.btor").unwrap();
    replace_anonymous_inputs_with_zero(&mut ctx, &mut sys);
    simplify_expressions(&mut ctx, &mut sys);
    // println!("{}", sys.serialize_to_str(&ctx));
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

/// Regression test for a bug that would prevent "_input_0" from being listed as a system input.
/// The reason was that we did not handle inputs that are also outputs well.
#[test]
fn parse_sha3_keccak_and_check_that_all_anonymous_inputs_are_there() {
    let (ctx, sys) =
        btor2::parse_file("../inputs/repair/sha3_keccak.w2.replace_variables.btor").unwrap();
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

/// Regression test to make sure that when an input that is also an output is set to zero,
/// it is maintained.
#[test]
fn parse_sha3_keccak_and_replace_anonymous_inputs() {
    let (mut ctx, mut sys) =
        btor2::parse_file("../inputs/repair/sha3_keccak.w2.replace_variables.btor").unwrap();
    replace_anonymous_inputs_with_zero(&mut ctx, &mut sys);
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

fn skip_first_line(value: &str) -> &str {
    match value.char_indices().find(|(_, c)| *c == '\n') {
        None => "",
        Some((char_ii, _)) => &value[(char_ii + 1)..],
    }
}

#[test]
#[ignore] // TODO: this test was randomly failing on the CI, but passing locally.
fn parse_lakeroad_dsp48_e2() {
    let (mut ctx, mut sys) = btor2::parse_file("../inputs/lakeroad/DSP48E2.btor").unwrap();
    replace_anonymous_inputs_with_zero(&mut ctx, &mut sys);
    simplify_expressions(&mut ctx, &mut sys);
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

const INVERTER: &str = r#"
; BTOR description generated by Yosys 0.51 (git sha1 c4b5190229616f7ebf8197f43990b4429de3e420, clang++ 16.0.0 -fPIC -O3) for module inverter_d0.
1 sort bitvec 32
2 input 1 a ; /Users/nikilshyamsunder/capra/protocols/examples/inverters/inverter_d0.v:2.19-2.20
3 not 1 2 $not$/Users/nikilshyamsunder/capra/protocols/examples/inverters/inverter_d0.v:6$1 ; /Users/nikilshyamsunder/capra/protocols/examples/inverters/inverter_d0.v:6.16-6.18
4 output 3 s ; /Users/nikilshyamsunder/capra/protocols/examples/inverters/inverter_d0.v:3.19-3.20
; end of yosys output
"#;

#[test]
fn serialize_inverter() {
    let mut ctx = Context::default();
    let sys = btor2::parse_str(&mut ctx, INVERTER, Some("inverter_d0")).unwrap();
    // check to make sure the output and input names are correct
    let inputs: Vec<_> = sys
        .inputs
        .iter()
        .map(|i| ctx.get_symbol_name(*i).unwrap().to_string())
        .collect();
    let outputs: Vec<_> = sys
        .outputs
        .iter()
        .map(|o| ctx[o.name].to_string())
        .collect();
    assert_eq!(inputs, ["a"]);
    assert_eq!(outputs, ["s"]);
    insta::assert_snapshot!(sys.serialize_to_str(&ctx));
}

#[test]
fn parse_llsdspi_bug_c1_instrumented() {
    let (mut ctx, mut sys) =
        btor2::parse_file("../inputs/repair/llsdspi_bug_c1.instrumented.btor").unwrap();
    replace_anonymous_inputs_with_zero(&mut ctx, &mut sys);
    simplify_expressions(&mut ctx, &mut sys);
    println!("{}", sys.serialize_to_str(&ctx));
}
