// Progressive MSA: k-mer distance + UPGMA guide tree + profile-profile Gotoh DP.
// Gap model: 1-piece affine g(l)=O+l*E, or 2-piece affine g(l)=min(O1+l*E1, O2+l*E2).
//
// Usage: msa <in_fasta> <mode:1|2> <O> <E>                    (1-piece)
//        msa <in_fasta> <mode:1|2> <O1> <E1> <O2> <E2>        (2-piece)
// Writes aligned FASTA to stdout. Writes "RUNTIME_MS=x PEAK_DP_BYTES=y" to stderr.

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

static const double NEG_INF = -1e18;
static const char BASES[4] = {'A','C','G','T'};
static double MATCH[4][4];

int base_idx(char c) {
    switch (c) { case 'A': return 0; case 'C': return 1; case 'G': return 2; case 'T': return 3; }
    return -1;
}

struct Cluster {
    vector<string> rows; // aligned rows for members currently in this cluster (all same length)
    vector<string> names;
    int size() const { return (int)rows.size(); }
    int len() const { return rows.empty() ? 0 : (int)rows[0].size(); }
};

vector<array<double,4>> profile_freqs(const Cluster &c) {
    int L = c.len(), n = c.size();
    vector<array<double,4>> freq(L, {0,0,0,0});
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < L; j++) {
            int b = base_idx(c.rows[i][j]);
            if (b >= 0) freq[j][b] += 1.0 / n;
        }
    }
    return freq;
}

struct GapParams {
    int npiece;
    double O1, E1, O2, E2;
};

long long g_peak_bytes = 0;

// Align two profiles, return list of (iA,iB) with -1 meaning gap on that side.
vector<pair<int,int>> gotoh_align(const vector<array<double,4>> &fa, const vector<array<double,4>> &fb,
                                   const GapParams &gp) {
    int LA = (int)fa.size(), LB = (int)fb.size();
    int R = LA+1, C = LB+1;
    auto idx = [&](int i,int j){ return i*C+j; };

    vector<double> M(R*C, NEG_INF);
    vector<double> Ix1(R*C, NEG_INF), Iy1(R*C, NEG_INF);
    vector<double> Ix2, Iy2;
    bool two = gp.npiece == 2;
    if (two) { Ix2.assign(R*C, NEG_INF); Iy2.assign(R*C, NEG_INF); }

    // traceback codes: 0=diag(M),1=up(Ix),2=left(Iy) for M matrix
    // for Ix: 0=open(from M),1=extend piece1,2=open piece2... encode piece via separate arrays
    vector<uint8_t> ptrM(R*C), ptrIx(R*C), ptrIy(R*C); // 0/1 meaning open/extend; piece stored implicitly by which array had max
    vector<uint8_t> pieceIx(R*C), pieceIy(R*C); // which piece (0 or 1) is currently best

    long long bytes = (long long)(R*C) * (sizeof(double)*(two?6:2) + sizeof(uint8_t)*4);
    g_peak_bytes = max(g_peak_bytes, bytes);

    M[idx(0,0)] = 0;
    for (int i = 1; i <= LA; i++) {
        double o1 = -gp.O1 - (i-1)*gp.E1;
        Ix1[idx(i,0)] = o1;
        pieceIx[idx(i,0)] = 0;
        if (two) {
            double o2 = -gp.O2 - (i-1)*gp.E2;
            Ix2[idx(i,0)] = o2;
            if (Ix2[idx(i,0)] > Ix1[idx(i,0)]) pieceIx[idx(i,0)] = 1;
        }
    }
    for (int j = 1; j <= LB; j++) {
        double o1 = -gp.O1 - (j-1)*gp.E1;
        Iy1[idx(0,j)] = o1;
        pieceIy[idx(0,j)] = 0;
        if (two) {
            double o2 = -gp.O2 - (j-1)*gp.E2;
            Iy2[idx(0,j)] = o2;
            if (Iy2[idx(0,j)] > Iy1[idx(0,j)]) pieceIy[idx(0,j)] = 1;
        }
    }

    for (int i = 1; i <= LA; i++) {
        for (int j = 1; j <= LB; j++) {
            // Ix (gap in B, consume A)
            double open1 = M[idx(i-1,j)] - gp.O1;
            double ext1 = Ix1[idx(i-1,j)] - gp.E1;
            double v1 = max(open1, ext1);
            Ix1[idx(i,j)] = v1;
            uint8_t bestpiece = 0; double bestv = v1;
            uint8_t op1 = (ext1 > open1) ? 1 : 0;
            if (two) {
                double open2 = M[idx(i-1,j)] - gp.O2;
                double ext2 = Ix2[idx(i-1,j)] - gp.E2;
                double v2 = max(open2, ext2);
                Ix2[idx(i,j)] = v2;
                if (v2 > bestv) { bestv = v2; bestpiece = 1; }
                if (bestpiece==1) ptrIx[idx(i,j)] = (ext2>open2)?1:0; else ptrIx[idx(i,j)] = op1;
            } else {
                ptrIx[idx(i,j)] = op1;
            }
            pieceIx[idx(i,j)] = bestpiece;

            // Iy (gap in A, consume B)
            double openy1 = M[idx(i,j-1)] - gp.O1;
            double exty1 = Iy1[idx(i,j-1)] - gp.E1;
            double vy1 = max(openy1, exty1);
            Iy1[idx(i,j)] = vy1;
            uint8_t bestpiecey = 0; double bestvy = vy1;
            uint8_t opy1 = (exty1 > openy1) ? 1 : 0;
            if (two) {
                double openy2 = M[idx(i,j-1)] - gp.O2;
                double exty2 = Iy2[idx(i,j-1)] - gp.E2;
                double vy2 = max(openy2, exty2);
                Iy2[idx(i,j)] = vy2;
                if (vy2 > bestvy) { bestvy = vy2; bestpiecey = 1; }
                if (bestpiecey==1) ptrIy[idx(i,j)] = (exty2>openy2)?1:0; else ptrIy[idx(i,j)] = opy1;
            } else {
                ptrIy[idx(i,j)] = opy1;
            }
            pieceIy[idx(i,j)] = bestpiecey;

            double ixv = bestv, iyv = bestvy;

            // score
            double s = 0;
            for (int a = 0; a < 4; a++) for (int b = 0; b < 4; b++)
                s += fa[i-1][a] * fb[j-1][b] * MATCH[a][b];

            double dM = M[idx(i-1,j-1)], dIx = Ix1[idx(i-1,j-1)], dIy = Iy1[idx(i-1,j-1)];
            if (two) { dIx = max(dIx, Ix2[idx(i-1,j-1)]); dIy = max(dIy, Iy2[idx(i-1,j-1)]); }
            double best = dM; uint8_t code = 0;
            if (dIx > best) { best = dIx; code = 1; }
            if (dIy > best) { best = dIy; code = 2; }
            M[idx(i,j)] = best + s;
            ptrM[idx(i,j)] = code;
        }
    }

    // pick best end state
    double endM = M[idx(LA,LB)];
    double endIx = Ix1[idx(LA,LB)]; if (two) endIx = max(endIx, Ix2[idx(LA,LB)]);
    double endIy = Iy1[idx(LA,LB)]; if (two) endIy = max(endIy, Iy2[idx(LA,LB)]);
    int state = 0; double best = endM;
    if (endIx > best) { best = endIx; state = 1; }
    if (endIy > best) { best = endIy; state = 2; }

    vector<pair<int,int>> aln;
    int i = LA, j = LB;
    while (i > 0 || j > 0) {
        if (i == 0) { aln.push_back({-1, j-1}); j--; continue; }
        if (j == 0) { aln.push_back({i-1, -1}); i--; continue; }
        if (state == 0) {
            aln.push_back({i-1, j-1});
            state = ptrM[idx(i,j)];
            i--; j--;
        } else if (state == 1) {
            aln.push_back({i-1, -1});
            state = ptrIx[idx(i,j)]; // 0 = came from M(open), 1 = extend (stay in Ix)
            i--;
        } else {
            aln.push_back({-1, j-1});
            state = ptrIy[idx(i,j)];
            j--;
        }
    }
    reverse(aln.begin(), aln.end());
    return aln;
}

Cluster merge_clusters(const Cluster &A, const Cluster &B, const GapParams &gp) {
    auto fa = profile_freqs(A);
    auto fb = profile_freqs(B);
    auto aln = gotoh_align(fa, fb, gp);
    Cluster out;
    int nA = A.size(), nB = B.size();
    for (int k = 0; k < nA; k++) out.rows.push_back(""), out.names.push_back(A.names[k]);
    for (int k = 0; k < nB; k++) out.rows.push_back(""), out.names.push_back(B.names[k]);
    for (auto &pr : aln) {
        int ia = pr.first, ib = pr.second;
        for (int k = 0; k < nA; k++) out.rows[k] += (ia>=0) ? A.rows[k][ia] : '-';
        for (int k = 0; k < nB; k++) out.rows[nA+k] += (ib>=0) ? B.rows[k][ib] : '-';
    }
    return out;
}

int main(int argc, char** argv) {
    if (argc < 5) {
        fprintf(stderr, "usage: msa <in_fasta> <mode:1|2> <O> <E> [<O2> <E2>]\n");
        return 1;
    }
    // match/mismatch matrix
    for (int a=0;a<4;a++) for(int b=0;b<4;b++) MATCH[a][b] = (a==b) ? 1.0 : -3.0;

    string in_fasta = argv[1];
    int mode = atoi(argv[2]);
    GapParams gp;
    gp.npiece = mode;
    gp.O1 = atof(argv[3]); gp.E1 = atof(argv[4]);
    if (mode == 2) {
        gp.O2 = atof(argv[5]); gp.E2 = atof(argv[6]);
    }

    // read fasta
    ifstream in(in_fasta);
    vector<string> names, seqs;
    string line, cur;
    while (getline(in, line)) {
        if (line.empty()) continue;
        if (line[0] == '>') {
            if (!cur.empty()) seqs.push_back(cur);
            names.push_back(line.substr(1));
            cur.clear();
        } else cur += line;
    }
    if (!cur.empty()) seqs.push_back(cur);
    int n = (int)seqs.size();

    auto t0 = chrono::high_resolution_clock::now();

    // k-mer (k=4) profile distance
    int K = 4;
    int dim = 1; for (int i=0;i<K;i++) dim*=4;
    vector<vector<double>> kmer(n, vector<double>(dim, 0.0));
    for (int s = 0; s < n; s++) {
        const string &seq = seqs[s];
        int L = (int)seq.size();
        for (int p = 0; p + K <= L; p++) {
            int code = 0; bool ok = true;
            for (int k = 0; k < K; k++) {
                int b = base_idx(seq[p+k]);
                if (b < 0) { ok = false; break; }
                code = code*4 + b;
            }
            if (ok) kmer[s][code] += 1;
        }
        double norm = 0; for (double v : kmer[s]) norm += v*v;
        norm = sqrt(norm);
        if (norm > 0) for (double &v : kmer[s]) v /= norm;
    }
    vector<vector<double>> dist(n, vector<double>(n, 0.0));
    for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) {
        double dot = 0; for (int k = 0; k < dim; k++) dot += kmer[i][k]*kmer[j][k];
        dist[i][j] = 1.0 - dot;
    }

    // UPGMA
    vector<int> cluster_size(n, 1);
    vector<int> alive_id(n); for (int i=0;i<n;i++) alive_id[i]=i;
    map<int, Cluster> clusters;
    for (int i = 0; i < n; i++) {
        Cluster c; c.rows = {seqs[i]}; c.names = {names[i]};
        clusters[i] = c;
    }
    map<int, map<int,double>> D;
    for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) if (i!=j) D[i][j] = dist[i][j];
    vector<int> live;
    for (int i = 0; i < n; i++) live.push_back(i);
    int next_id = n;

    while (live.size() > 1) {
        double best = 1e18; int bi=-1, bj=-1;
        for (size_t a = 0; a < live.size(); a++)
            for (size_t b = a+1; b < live.size(); b++) {
                int x = live[a], y = live[b];
                double d = D[x][y];
                if (d < best) { best = d; bi = x; bj = y; }
            }
        Cluster merged = merge_clusters(clusters[bi], clusters[bj], gp);
        int nid = next_id++;
        clusters[nid] = merged;
        live.erase(remove(live.begin(), live.end(), bi), live.end());
        live.erase(remove(live.begin(), live.end(), bj), live.end());
        // update distances (UPGMA average, weighted by member counts)
        int size_bi = clusters[bi].size(), size_bj = clusters[bj].size();
        for (int x : live) {
            double d = (D[bi][x]*size_bi + D[bj][x]*size_bj) / (size_bi+size_bj);
            D[nid][x] = d; D[x][nid] = d;
        }
        live.push_back(nid);
    }

    Cluster final = clusters[live[0]];

    auto t1 = chrono::high_resolution_clock::now();
    double ms = chrono::duration<double, milli>(t1-t0).count();

    for (int i = 0; i < (int)final.rows.size(); i++) {
        cout << ">" << final.names[i] << "\n" << final.rows[i] << "\n";
    }
    fprintf(stderr, "RUNTIME_MS=%.3f PEAK_DP_BYTES=%lld\n", ms, g_peak_bytes);
    return 0;
}