// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// Original author: Louis-Emile Ploix
// SPDX-License-Identifier: Apache-2.0

// aig-manip is a utility program for manipulating aiger files for our usecase.
// It supports 2 main functions, with a couple of extra utilities too:
// 1. selection, where we take an aiger file and delete asserts or transform them into assumes.
// 2. simulation, where take an AIW (an aiger witness file) and produce a VCD.
// Additionally we can:
// 1. Do COI analysis (i.e. what are all the named signals this wire depends on)
// 2. Do why queries, (i.e. whare are all the named signals that made this wire have the value it has)
// Some background on aiger (.aig) files:
// - They are essentially arrays of AND gates, which are represented by two integers.
//   The lowest bit of the number is 1 if the input should be considered negated, and 0 otherwise.
//   The remaining bits are the index either into the table of inputs, FFs or AND gates.
// - The AND gates are topologically sorted, so simulation is really easy.
// - There is no representation of assertions/assumptions etc, but through comments
// - Yosys provides us with the indices of wires through vmap and ywmap files.
//   vmap files list the indices of every wire.
//   ywmap files (are JSON) and list the indices of every assertion.

use std::collections::HashSet;
use std::{
    collections::HashMap,
    io::{BufReader, BufWriter},
};

use clap::Parser;

use crate::aiw::AigerWitness;
use crate::bitvec::BitVec;
use crate::vmap::VMap;
use crate::ywmap::YWMap;

mod aiw;
mod bitvec;
mod vmap;
mod ywmap;

#[derive(clap::Parser, Debug)]
struct Args {
    #[arg(help = "Path to .aig file to read.")]
    aig: String,

    #[clap(subcommand)]
    sub: Subcommand,
}

#[derive(clap::Subcommand, Debug)]
enum Subcommand {
    #[command(
        about = "Produce a specialised version of a given .aig file, suitable to run a modelchecker against directly."
    )]
    Select {
        #[arg(long, help = "Instead of converting asserts to assumes, delete them.")]
        no_assumes: bool,
        #[arg(help = "Path of .aig file to write output to")]
        out: String,
        #[arg(help = "Path to .ywmap file generated by yosys for this .aig")]
        yw: String,
        #[arg(help = "The step number we are constructing for.")]
        step: usize,
        #[arg(
            help = "The properties (in the given step) to keep, the rest will be deleted, unless there are no properties then they are all kept."
        )]
        prop: Vec<String>,
    },
    #[command(
        about = "Transform an aiger witness (.aiw) into a VCD file. Note that if child.aig was passed to the modelchecker, but parent.aig is it's parent, then use parent.aig."
    )]
    Simulate {
        #[arg(help = "Path to .ywmap file generated by yosys for this .aig")]
        yw: String,
        #[arg(help = "Path to .vmap file generated by yosys for this .aig")]
        vmap: String,
        #[arg(help = "Path to .aiw file")]
        trace: String,
        #[arg(help = "Path to write .vcd")]
        vcd: String,
    },
    #[command(about = "Produce a list of signals that influence a given wire")]
    Coi {
        #[arg(help = "Path to .vmap file generated by yosys for this .aig")]
        vmap: String,
        #[arg(help = "Full name of query signal")]
        name: String,
        #[arg(
            long,
            help = "Path to .ywmap file generated by yosys for this .aig (necessary if signal in question is an assertion)"
        )]
        ywmap: Option<String>,
    },
    #[command(about = "Produce a list of signals that influenced a given wire in a trace")]
    Why {
        #[arg(help = "Path to .vmap file generated by yosys for this .aig")]
        vmap: String,
        #[arg(help = "Path to .aiw file")]
        trace: String,
        #[arg(help = "Time step (starting from reset at 0) to query at")]
        step: usize,
        #[arg(help = "Full name of query signal")]
        signal: String,
        #[arg(
            long,
            short,
            help = "Bit of query signal of interest, otherwise all of them"
        )]
        bit: Option<usize>,
    },
}

/// Select properties of interest in the aig file
fn select(
    mut aig: aig::Aig,
    out: String,
    yw: String,
    step: usize,
    prop: Vec<String>,
    no_assumes: bool,
) {
    let ywmap = YWMap::parse(BufReader::new(
        std::fs::File::open(&yw).expect("Could not open YW file"),
    ));

    assert_eq!(aig.justice.len(), 0);
    assert_eq!(aig.fairness.len(), 0);

    assert_eq!(aig.bads.len(), ywmap.asserts.len());

    let mut bn = 0; // Index in list of names
    let mut b = 0; // Index in what remains of the list of bads in the AIG
    while b < aig.bads.len() {
        let name = &ywmap.asserts[bn][0][6..]
            .split_once('$')
            .expect("Expected \\top.*$n")
            .0;
        let (stepn, name) = name.split_once('_').expect("badly formed property");
        assert!(stepn.starts_with("Step"));
        let pstep: usize = stepn[4..].parse().expect("badly formed property name");
        let is_cover = name.ends_with("_Cover");
        bn += 1;

        if pstep == step && (prop.is_empty() || prop.iter().any(|x| x == name)) {
            if is_cover {
                // We encode covers as asserts since otherwise they are lost before going into the AIG file,
                // this means we need to negate them if we want to do BMC for them.
                println!("Keeping {name} (cover statement, negating)");
                aig.bads[b] = !aig.bads[b];
            } else {
                println!("Keeping {name}");
            }
            b += 1;
        } else if no_assumes || is_cover || pstep >= step {
            println!("Deleting {}", name);
            aig.bads.remove(b);
        } else if pstep < step {
            println!("Assert2Assume {}", name);
            aig.constraints.push(!aig.bads[b]);
            aig.bads.remove(b);
        } else {
            unreachable!()
        }
    }

    aig.to_file(&out, false);
}

/// Simulate a trace and produce a VCD file
fn simulate(aig: aig::Aig, yw: String, vmap: String, trace: String, out: String) {
    let ywmap = YWMap::parse(BufReader::new(
        std::fs::File::open(&yw).expect("Could not open YW file"),
    ));
    let mut vmap = VMap::parse(BufReader::new(
        std::fs::File::open(vmap).expect("Could not open vmap"),
    ));
    vmap.wires.extend(ywmap.vmap_wires(&aig));

    let mut vcd = vcd::Writer::new(BufWriter::new(
        std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(out)
            .expect("Could not open VCD out"),
    ));
    vcd.timescale(1, vcd::TimescaleUnit::US).unwrap(); // This is arbitrary

    let hier = vmap.to_hierarchy();
    let mut extract = Vec::new();
    vcd.add_module("top").unwrap();
    hier.append_to_vcd(&mut vcd, &mut extract)
        .expect("Could not write wire hierarchy to VCD");
    vcd.upscope().unwrap();
    vcd.enddefinitions().unwrap();

    let witness = AigerWitness::parse(BufReader::new(
        std::fs::File::open(&trace).expect("could not open witness file"),
    ))
    .expect("error reading witness file");

    let mut sim = witness.init_simulation(&aig);
    for (s, step) in witness.inputs.iter().enumerate() {
        sim.simulate(step.iter().map(|x| x.into()).collect());
        println!("Step {}", s);

        for (a, [assert]) in ywmap.asserts.iter().enumerate() {
            if assert.contains("_Cover") {
                continue;
            }
            if sim.value(aig.bads[a]).is_true() {
                println!("Violated assertion {assert}");
            }
        }
        for (a, [assume]) in ywmap.assumes.iter().enumerate() {
            if sim.value(aig.constraints[a]).is_false() {
                println!("Violated assumption {assume}");
            }
        }

        vcd.timestamp(s as u64).unwrap();
        for instr in &extract {
            let bits: Vec<_> = instr
                .edges
                .iter()
                .map(|edge| {
                    let v = sim.value(*edge);
                    if v.is_true() {
                        vcd::Value::V1
                    } else if v.is_false() {
                        vcd::Value::V0
                    } else {
                        vcd::Value::X
                    }
                })
                .collect();
            // FIXME: We don't actually need to do this if the values haven't changed.
            if let [bit] = &bits[..] {
                vcd.change_scalar(instr.idcode, *bit).unwrap();
            } else {
                vcd.change_vector(instr.idcode, bits).unwrap();
            }
        }
    }
    vcd.timestamp(witness.inputs.len() as u64).unwrap();
}

/// Determine the COI for a wire
fn coi(aig: aig::Aig, vmap: String, signal: String, ywmap: Option<String>) {
    let mut vmap = VMap::parse(BufReader::new(
        std::fs::File::open(vmap).expect("Could not open vmap"),
    ));
    if let Some(ywmap) = ywmap {
        let ywmap = YWMap::parse(BufReader::new(
            std::fs::File::open(&ywmap).expect("Could not open YW file"),
        ));
        vmap.wires.extend(ywmap.vmap_wires(&aig));
    }
    let hier = vmap.to_hierarchy();
    let group = hier.find(&signal.split(".").collect::<Vec<_>>());
    let Some(group) = group else {
        println!("Name not found");
        return;
    };

    let named_wires = hier.named_aiger_vars();
    let mut found = HashSet::new();
    let mut visited = HashSet::new();

    // Straightforward DFS on each bit to find the wires we care about,
    // stopping as soon as a wire is named
    for bit in &group.bits {
        let node = &aig.nodes[bit.index >> 1];
        if !node.is_and() {
            continue;
        }
        let mut to_visit = vec![node.fanin0().node_id(), node.fanin1().node_id()];
        while let Some(nodeid) = to_visit.pop() {
            if let Some(wire) = named_wires.get(&nodeid) {
                found.insert((wire.path, node.is_and()));
                continue;
            }
            let node = &aig.nodes[nodeid];
            if !node.is_and() {
                continue;
            }
            if visited.insert(node.fanin0()) {
                to_visit.push(node.fanin0().node_id());
            }
            if visited.insert(node.fanin1()) {
                to_visit.push(node.fanin1().node_id());
            }
        }
    }

    for (found, and) in found {
        if and {
            println!("{} (and node)", found);
        } else {
            println!("{}", found);
        }
    }
}

/// Determine the why signals for a wire in a trace at a given point in time
fn why(
    aig: aig::Aig,
    vmap: String,
    trace: String,
    step: usize,
    signal: String,
    bit: Option<usize>,
) {
    let vmap = VMap::parse(BufReader::new(
        std::fs::File::open(vmap).expect("Could not open vmap"),
    ));

    let hier = vmap.to_hierarchy();
    let group = hier.find(&signal.split(".").collect::<Vec<_>>());
    let Some(group) = group else {
        println!("Name not found");
        return;
    };

    let witness = AigerWitness::parse(BufReader::new(
        std::fs::File::open(&trace).expect("could not open witness file"),
    ))
    .expect("error reading witness file");
    let sim = witness.simulate_to_step(&aig, step);

    let named_wires = hier.named_aiger_vars();
    let mut found = HashMap::new();
    let mut visited = HashSet::new();

    let bits = match bit {
        None => &group.bits[..],
        Some(i) => &[group.bits[i]],
    };
    // Similar DFS as above, but now implementing tri-state logic to
    // decide where to explore into.
    for bit in bits {
        let node = &aig.nodes[bit.index >> 1];
        if !node.is_and() {
            continue;
        }
        let mut to_visit = vec![
            (node.fanin0().node_id(), vec![bit.index >> 1]),
            (node.fanin1().node_id(), vec![bit.index >> 1]),
        ];
        while let Some((nodeid, path)) = to_visit.pop() {
            if let Some(wire) = named_wires.get(&nodeid) {
                found.insert(nodeid, (wire.path, wire));
                continue;
            }
            let node = &aig.nodes[nodeid];
            if !node.is_and() {
                continue;
            }
            let left = sim.value(node.fanin0());
            let right = sim.value(node.fanin1());
            assert!(!left.is_none() && !right.is_none());

            let mut new_path = path.clone();
            new_path.push(nodeid);
            match (left.is_true(), right.is_true()) {
                // If a = b, then both are response for the value of a & b. Arguably neither are but that's not helpful.
                (true, true) | (false, false) => {
                    if visited.insert(node.fanin0().node_id()) {
                        to_visit.push((node.fanin0().node_id(), new_path.clone()));
                    }
                    if visited.insert(node.fanin1().node_id()) {
                        to_visit.push((node.fanin1().node_id(), new_path));
                    }
                }
                // true & false = false, and the false is responsible
                (true, false) => {
                    if visited.insert(node.fanin1().node_id()) {
                        to_visit.push((node.fanin1().node_id(), new_path));
                    }
                }
                // false & true = false, and the false is responsible
                (false, true) => {
                    if visited.insert(node.fanin0().node_id()) {
                        to_visit.push((node.fanin0().node_id(), new_path));
                    }
                }
            }
        }
    }

    let mut lines = Vec::new();
    for (wirei, (_path, group)) in found.iter() {
        let wire = group.bits.iter().find(|x| x.index >> 1 == *wirei).unwrap();
        let v = sim.value(wire.aig_edge());
        let c = if v.is_true() {
            '1'
        } else if v.is_false() {
            '0'
        } else {
            'x'
        };
        lines.push(format!(
            "{} [{}] = {c}     ({})",
            wire.path, wire.offset, wire.index,
        ));
    }
    lines.sort();
    for line in lines {
        println!("{}", line);
    }
}

/// Check the AIG file is topologically sorted.
pub fn check_aig(aig: &aig::Aig) {
    let mut ok = BitVec::fill_zero(aig.nodes.len());
    ok.set(0, true);
    for i in 0..aig.inputs.len() {
        ok.set(aig.inputs[i], true);
    }
    for i in 0..aig.latchs.len() {
        ok.set(aig.latchs[i].input, true);
    }
    for i in aig.nodes_range() {
        if aig.nodes[i].is_and() {
            assert!(ok.get(aig.nodes[i].fanin0().node_id()).unwrap());
            assert!(ok.get(aig.nodes[i].fanin1().node_id()).unwrap());
            ok.set(i, true);
        }
    }
    for i in 0..aig.latchs.len() {
        assert!(ok.get(aig.latchs[i].next.node_id()).unwrap());
    }
}

fn main() {
    let args = Args::parse();
    let aig = aig::Aig::from_file(&args.aig);
    check_aig(&aig); // Never actually failed! Yosys maintains this invariant

    match args.sub {
        Subcommand::Select {
            out,
            yw,
            step,
            prop,
            no_assumes,
        } => select(aig, out, yw, step, prop, no_assumes),
        Subcommand::Simulate {
            yw,
            vmap,
            trace,
            vcd: out,
        } => simulate(aig, yw, vmap, trace, out),
        Subcommand::Coi { vmap, name, ywmap } => coi(aig, vmap, name, ywmap),
        Subcommand::Why {
            vmap,
            trace,
            step,
            signal,
            bit,
        } => why(aig, vmap, trace, step, signal, bit),
    }
}
