use dep::std;


// tx circuit + publisher
fn main(
  tx_in: pub [Field; 16], // Publicly posted txs for forced inclusion
  amount_public_in: pub Field,
  amount_public_out: pub Field,
  commitment_out: pub [Field; 16], // tx[0] -- [...h(utxo)] --> ONLY CONTAINS FIRST TX COMMITMENTS
  recipient: pub Field,
  oracle: pub Field,
  old_root: pub Field, // [Publisher] Feed in from contract
  new_root: pub Field, // Calculate and Return
  nullifier_hashes: pub [Field; 16],
  secrets: [Field; 16],
  utxo_in: [Field; 48], // 16 * [owner, amount, asset_type]
  utxo_out: [Field; 48], // 16 * [owner, amount, asset_type]
  roots: [Field; 64], // 16 * [utxo_root, tx_root, batch_root, historic_root]
  leaves: [Field; 64], // 16 * [utxo_leaf, tx_leaf, batch_leaf, historic_leaf]
  indexes: [Field; 64], // 16 * [utxo_index, tx_index, batch_index, historic_index]
  hash_path: [Field; 288], // 16 * [utxo_path, tx_path, batch_path, historic_path]
) { // -> [Field; 16] {

    let trees: Field = 4; // UTXO->Tx->Batch->HistoricalState->CurrentState

    // Initialize input and output tallies as public amounts
    let mut sum_in: Field = amount_public_in;
    let mut sum_out: Field = amount_public_out;

    // Iterate over inputs
    for i in 0..16 {

        if (utxo_in[i*3 + 1] != 0) {

            // Assert h(secret) == owner - make sure user can spend this utxo
            let owner = std::hash::pedersen([secrets[i]]);
            assert(owner[0] == utxo_in[i*3 + 0]);

            // Nullifier is h(secret, secret) to avoid leaking spender - maybe use h(comm, secret) or something else later
            assert(nullifier_hashes[i] == std::hash::pedersen([secrets[i], secrets[i]])[0]);

            let commitment_in = std::hash::pedersen([utxo_in[i*3 + 0], utxo_in[i*3 + 1], utxo_in[i*3 + 2]])[0];

            // Initialize hash path arrays
            let mut hash_path_utxo: [Field; 4] = [0; 4]; // UTXO->Tx Limit
            let mut hash_path_tx: [Field; 4] = [0; 4]; // Tx->Batch Limit
            let mut hash_path_batch: [Field; 5] = [0; 5]; // Batch->State Limit
            let mut hash_path_historic: [Field; 5] = [0; 5]; // Total State Limit

            for j in 0..4 { // 4 levels per hash path (for now)
                hash_path_utxo[j] = hash_path[18*i + 0 + j];
                hash_path_tx[j] = hash_path[18*i + 4 + j];
            }

            for l in 0..5 { // 5 levels for batch and historic trees (for now)
                hash_path_batch[l] = hash_path[18*i + 8 + l];
                hash_path_historic[l] = hash_path[18*i + 13 + l];
            }

            let leaf_tx = leaves[trees * i + 1]; 
            let leaf_batch = leaves[trees * i + 2];
            let leaf_historic = leaves[trees * i + 3];

            let index_utxo = indexes[trees * i + 0];
            let index_tx = indexes[trees * i + 1];
            let index_batch = indexes[trees * i + 2];
            let index_historic = indexes[trees * i + 3];

            // h([...utxo]) == root_utxo / tx_id
            // h([...tx]) == root_batch
            // h([...root_batch]) == current_root
            // h([...root_leaves]) == historic_root

            // leaf_batch = h(root_tx, oracle)
            // leaf_historic = h()

            //let root_utxo = roots[trees * i + 0];
            let root_tx = roots[trees * i + 1];
            let root_batch = roots[trees * i + 2];
            let root_historic = roots[trees * i + 3];

            // fn compute_merkle_root(
                // leaf : Field, 
                // index : Field,
                // hash_path: [Field]
                // ) -> Field

            // utxo_root == tx_id
            let utxo_root = std::merkle::compute_merkle_root(
                commitment_in,
                index_utxo,
                hash_path_utxo
                );
            assert(utxo_root == leaf_tx);

            let tx_root = std::merkle::compute_merkle_root(
                leaf_tx,
                index_tx,
                hash_path_tx
                );
            assert(tx_root == root_tx);

            let batch_root = std::merkle::compute_merkle_root(
                leaf_batch, 
                index_batch,
                hash_path_batch
                );
            assert(batch_root == root_batch);

            let historic_root = std::merkle::compute_merkle_root(
                leaf_historic, 
                index_historic,
                hash_path_historic
                );
            assert(historic_root == root_historic);
            
            // Add utxo value to input total
            sum_in += utxo_in[i*3 + 1];
        }
    }

    // Iterate over outputs
    for k in 0..16 {
        if (utxo_out[k*3 + 1] != 0) {
            let commitment_out_calc = std::hash::pedersen([utxo_out[k*3 + 0], utxo_out[k*3 + 1], utxo_out[k*3 + 2]]);
            // Constraint check here to prevent loss of funds
            assert(commitment_out_calc[0] == commitment_out[k]);
            sum_out += utxo_out[k*3 + 1];
        } 
        else {
            let zero_hash = 0xf35fcb490b7ea67c3ac26ed530fa5d8dfe8be344e7177ebb63fe02723fb6f725 as Field;
            assert(commitment_out[k] == zero_hash);
        }
    }

    // [Publisher] Construct the new root
    // - newRoot->h(oldRoot, h(h([...h([...utxo])]), h(oracleData)))
    // - batchRoot->h(batch, h(oracleData))
    // - batch->h([...tx]
    // - tx->h([...utxo])

    // Find Tx aka UTXORoot (UTXO->Tx)]
    let utxo_root_calc: Field = pedersen_tree_four(commitment_out);
    assert(tx_in[0] == utxo_root_calc); // TODO: Check all indices
    let tx_root_calc: Field = pedersen_tree_four(tx_in); // Only works for MAX_BATCH_SIZE = 16
    assert(oracle == std::hash::pedersen([0])[0]); // TODO: Assert the actual data matches the assets contract
    let batch_root_calc: Field = std::hash::pedersen([tx_root_calc, oracle])[0];
    let new_root_calc: Field = std::hash::pedersen([batch_root_calc, old_root])[0];
    assert(new_root == new_root_calc); // Make sure we submitted the correct h(old_root, batch)

    // root = pedersen(Batch, Oracle) -- Oracle is zero_values[0] for now
    // newRoot = pedersen(root, oldRoot)
    
    // Check the final sums
    assert(sum_in == sum_out);

    // Silence the warning. The compiler probably optimizes this out anyway
    assert(recipient == recipient);


}

fn pedersen_tree_four(leaves: [Field; 16]) -> Field {
    //let mut num_hashes: u8 = 8; // make this leaves.len() / 2 ?
    let mut tx_tree: [Field; 16] = leaves; // make this length num_hashes ?
    for l in 0..8 {
        tx_tree[l] = std::hash::pedersen([tx_tree[2*l], tx_tree[2*l + 1]])[0];
    }
    for l in 0..4 {
        tx_tree[l] = std::hash::pedersen([tx_tree[2*l], tx_tree[2*l + 1]])[0];
    }
    for l in 0..2 {
        tx_tree[l] = std::hash::pedersen([tx_tree[2*l], tx_tree[2*l + 1]])[0];
    }
    std::hash::pedersen([tx_tree[0], tx_tree[1]])[0] // root
}