// Simulate sequence family evolution along a random binary tree with K2P substitutions
// and mixed-geometric indels. Emits unaligned FASTA + ground-truth alignment FASTA per family.
//
// Usage: simulate <seed> <pi_long> <n_families> <n_taxa> <root_len> <out_prefix>
// Writes: <out_prefix>_famK_unaligned.fa , <out_prefix>_famK_true.fa  for K=0..n_families-1

#include <bits/stdc++.h>
using namespace std;

static mt19937 rng;

double runif(double a, double b) {
    uniform_real_distribution<double> d(a, b);
    return d(rng);
}
int rpoisson(double lam) {
    if (lam <= 0) return 0;
    poisson_distribution<int> d(lam);
    return d(rng);
}
int rgeom_len(double mean) {
    // geometric length >=1 with given mean
    double p = 1.0 / mean;
    geometric_distribution<int> d(p);
    return d(rng) + 1;
}
char rand_base() {
    const char b[4] = {'A','C','G','T'};
    uniform_int_distribution<int> d(0,3);
    return b[d(rng)];
}
int base_idx(char c) {
    switch(c) { case 'A': return 0; case 'C': return 1; case 'G': return 2; case 'T': return 3; }
    return 0;
}
char mutate_k2p(char c) {
    // transition pairs: A<->G , C<->T ; ts:tv = 2:1 total => P(transition)=2/3
    int idx = base_idx(c);
    double r = runif(0,1);
    if (r < 2.0/3.0) {
        // transition
        static const int trans[4] = {2,3,0,1}; // A->G, C->T, G->A, T->C
        return "ACGT"[trans[idx]];
    } else {
        // transversion: two possible targets, pick uniformly
        static const int tv[4][2] = {{1,3},{0,2},{1,3},{0,2}}; // A->{C,T}, C->{A,G}, G->{C,T}, T->{A,G}
        int which = (runif(0,1) < 0.5) ? 0 : 1;
        return "ACGT"[tv[idx][which]];
    }
}

struct Node {
    int left = -1, right = -1; // children indices, -1 if leaf
    double blen_left = 0, blen_right = 0;
    int leaf_id = -1;
};

// live representation carried down the tree
struct LiveSeq {
    vector<int> ids;    // master column ids, in order
    vector<char> bases; // current base per id
};

struct Sim {
    int next_id = 0;
    list<int> master;
    unordered_map<int, list<int>::iterator> pos;

    int new_root_column() {
        int id = next_id++;
        master.push_back(id);
        pos[id] = prev(master.end());
        return id;
    }
};

void apply_branch(Sim &sim, const LiveSeq &parent, double blen, double indel_rate, double pi_long,
                   LiveSeq &child) {
    child = parent;
    int len = (int)child.ids.size();
    // substitutions
    int nsub = rpoisson(blen * max(len,1) * 1.0);
    for (int t = 0; t < nsub; t++) {
        if (child.bases.empty()) break;
        uniform_int_distribution<int> d(0, (int)child.bases.size()-1);
        int p = d(rng);
        child.bases[p] = mutate_k2p(child.bases[p]);
    }
    // indels
    int nindel = rpoisson(blen * max(len,1) * indel_rate);
    for (int t = 0; t < nindel; t++) {
        int L = (int)child.ids.size();
        bool do_ins = (runif(0,1) < 0.5) || (L == 0);
        if (do_ins) {
            int ilen = (runif(0,1) < pi_long) ? rgeom_len(30.0) : rgeom_len(2.0);
            uniform_int_distribution<int> d(0, L); // insertion slot
            int slot = d(rng);
            // find master iterator to insert after: the id at position slot-1 in child.ids (if exists)
            list<int>::iterator insert_after;
            bool have_after = (slot > 0);
            if (have_after) insert_after = sim.pos[child.ids[slot-1]];
            vector<int> new_ids(ilen);
            vector<char> new_bases(ilen);
            list<int>::iterator cursor = have_after ? next(insert_after) : sim.master.begin();
            for (int k = 0; k < ilen; k++) {
                int id = sim.next_id++;
                cursor = sim.master.insert(cursor, id);
                sim.pos[id] = cursor;
                ++cursor;
                new_ids[k] = id;
                new_bases[k] = rand_base();
            }
            child.ids.insert(child.ids.begin()+slot, new_ids.begin(), new_ids.end());
            child.bases.insert(child.bases.begin()+slot, new_bases.begin(), new_bases.end());
        } else {
            int dlen = (runif(0,1) < pi_long) ? rgeom_len(30.0) : rgeom_len(2.0);
            dlen = min(dlen, L);
            if (dlen <= 0) continue;
            uniform_int_distribution<int> d(0, L - dlen);
            int start = d(rng);
            child.ids.erase(child.ids.begin()+start, child.ids.begin()+start+dlen);
            child.bases.erase(child.bases.begin()+start, child.bases.begin()+start+dlen);
        }
    }
}

int main(int argc, char** argv) {
    if (argc < 7) {
        fprintf(stderr, "usage: simulate <seed> <pi_long> <n_families> <n_taxa> <root_len> <out_prefix>\n");
        return 1;
    }
    unsigned seed = (unsigned)atoi(argv[1]);
    double pi_long = atof(argv[2]);
    int n_families = atoi(argv[3]);
    int n_taxa = atoi(argv[4]);
    int root_len = atoi(argv[5]);
    string out_prefix = argv[6];
    double indel_rate = 0.06; // events per site per unit branch length

    rng.seed(seed);

    for (int fam = 0; fam < n_families; fam++) {
        Sim sim;
        // random root sequence
        LiveSeq root;
        for (int i = 0; i < root_len; i++) {
            int id = sim.new_root_column();
            root.ids.push_back(id);
            root.bases.push_back(rand_base());
        }
        // build random binary tree topology over n_taxa leaves (caterpillar-ish random join order)
        vector<int> active_nodes;
        vector<Node> nodes;
        vector<LiveSeq> node_seq; // filled once computed top-down, index aligned with nodes vector for internal nodes
        // We'll build topology bottom-up (random joining like simple coalescent), storing tree as
        // a binary structure with an explicit root, then simulate top-down via recursion.
        struct TNode { int id; TNode* l=nullptr; TNode* r=nullptr; double bl=0,br=0; };
        vector<TNode*> pool;
        for (int i = 0; i < n_taxa; i++) pool.push_back(new TNode{i,nullptr,nullptr,0,0});
        int internal_id = n_taxa;
        while (pool.size() > 1) {
            uniform_int_distribution<int> d(0,(int)pool.size()-1);
            int a = d(rng);
            int b;
            do { b = d(rng); } while (b == a);
            TNode* na = pool[a]; TNode* nb = pool[b];
            TNode* parent = new TNode{internal_id++, na, nb, runif(0.01,0.05), runif(0.01,0.05)};
            vector<TNode*> newpool;
            for (int i = 0; i < (int)pool.size(); i++) if (i!=a && i!=b) newpool.push_back(pool[i]);
            newpool.push_back(parent);
            pool = newpool;
        }
        TNode* troot = pool[0];

        vector<string> leaf_unaligned(n_taxa), leaf_names(n_taxa);
        unordered_map<int, LiveSeq> leaf_live;

        function<void(TNode*, const LiveSeq&)> rec = [&](TNode* node, const LiveSeq& parent_seq) {
            if (!node->l && !node->r) {
                leaf_live[node->id] = parent_seq;
                return;
            }
            LiveSeq lc, rc;
            apply_branch(sim, parent_seq, node->bl, indel_rate, pi_long, lc);
            apply_branch(sim, parent_seq, node->br, indel_rate, pi_long, rc);
            rec(node->l, lc);
            rec(node->r, rc);
        };
        rec(troot, root);

        // build master column index map
        unordered_map<int,int> col_index;
        {
            int idx = 0;
            for (int id : sim.master) col_index[id] = idx++;
        }
        int total_cols = (int)sim.master.size();

        // write unaligned + true alignment FASTA
        string fn_un = out_prefix + "_fam" + to_string(fam) + "_unaligned.fa";
        string fn_tr = out_prefix + "_fam" + to_string(fam) + "_true.fa";
        ofstream fu(fn_un), ft(fn_tr);
        for (int t = 0; t < n_taxa; t++) {
            LiveSeq &ls = leaf_live[t];
            string unaligned(ls.bases.begin(), ls.bases.end());
            string aligned(total_cols, '-');
            for (size_t k = 0; k < ls.ids.size(); k++) {
                aligned[col_index[ls.ids[k]]] = ls.bases[k];
            }
            fu << ">taxon" << t << "\n" << unaligned << "\n";
            ft << ">taxon" << t << "\n" << aligned << "\n";
        }
        fu.close(); ft.close();

        for (auto p : pool) {} // no-op
    }
    fprintf(stderr, "simulate: seed=%u pi_long=%.2f n_families=%d n_taxa=%d root_len=%d -> %s_fam*.fa\n",
            seed, pi_long, n_families, n_taxa, root_len, out_prefix.c_str());
    return 0;
}