/*
  Copyright (C) 2004-2006,2009 Benjamin Redelings

  This file is part of BAli-Phy.

  BAli-Phy is free software; you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free
  Software Foundation; either version 2, or (at your option) any later
  version.

  BAli-Phy is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with BAli-Phy; see the file COPYING.  If not see
  <http://www.gnu.org/licenses/>.  */

#include "distance-methods.H"

#include <iostream>
#include <list>
#include <valarray>
#include "boost/tuple/tuple.hpp"
#include "inverse.H"

using std::list;
using std::vector;

using boost::dynamic_bitset;
using namespace boost::tuples;

/// Compute the sum of all path lengths delta[b] passing through each branch b.
vector<double> FastMTM(const Tree& T,const Matrix& D,
                       const vector<vector<int> >& leaf_sets) 
{
    const int N = T.n_leaves();

    vector<double> d(T.n_branches(),-1);

    // compute delta[i]*D for leaf taxa [i]
    for(int i=0;i<T.n_leafbranches();i++) {
        double sum = 0;
        for(int j=0;j<N;j++)
            if (j!=i)
                sum += D(i,j);
        d[i] = sum;
    }

    vector<const_branchview> branches = branches_from_node(T,T.n_leaves());
    std::reverse(branches.begin(), branches.end());

    for(int b=0;b<branches.size();b++) 
    {
        if (branches[b].is_leaf_branch()) continue;

        vector<const_branchview> children;
        append(branches[b].branches_after(),children);
    
        double temp = 0;
        for(int i=0;i<children.size();i++)
            temp += d[ children[i].undirected_name() ];

        double sum = 0;
        for(int i=0;i<children.size();i++)
            for(int j=0;j<i;j++) {
      
                const vector<int>& xset = leaf_sets[children[i]];
                const vector<int>& yset = leaf_sets[children[j]];

                for(int x=0;x<xset.size();x++)
                    for(int y=0;y<yset.size();y++)
                        sum += D(xset[x],yset[y]);
            }

        d[branches[b].undirected_name()] = temp - 2.0*sum;
    }
  
    return d;
}

vector<double> SlowMTM(const Tree& T,const Matrix& D,
                       const vector<vector<int> >& leaf_sets) 
{
    vector<double> d(T.n_branches(),0);

    for(int b=0;b<T.n_branches();b++) {
        vector<int> xset = leaf_sets[T.directed_branch(b)];
        vector<int> yset = leaf_sets[T.directed_branch(b).reverse()];

        for(int i=0;i<xset.size();i++)
            for(int j=0;j<yset.size();j++)
                d[b] += D(xset[i],yset[j]);
    }
    return d;
}

// Solve (A^t * W * A) b = (A^t * W * d) , where W is "diagonal" (on paths)

// M1 = sum_ij (A[k,ij] * A[l,ij] * W[ij])  (symmetric)

// M2 = sum_ij (A[k,ij] * W[ij] * D[ij])

// M1[k,l] * b[l] = M2[k]

bool A(const Tree& T,int k,int i,int j)
{
    const dynamic_bitset<>& partition = T.partition(k);
    return partition[i] != partition[j];
}

vector<double> LeastSquares(const Tree& T, const Matrix & D, const Matrix& W,
                            const vector<vector<int> >& leaf_sets) 
{
    const int N = T.n_leaves();
    const int B = T.n_branches();

    assert(D.size1() == N);
    assert(D.size2() == N);
    assert(W.size1() == N);
    assert(W.size2() == N);
    Matrix WD = D;
    for(int i=0;i<WD.size1();i++)
        for(int j=0;j<WD.size2();j++)
            WD(i,j) *= W(i,j);

    vector<double> M2v = FastMTM(T,WD,leaf_sets);
    Matrix M2(B,1);
    for(int i=0;i<B;i++)
        M2(i,0) = M2v[i];

    Matrix M1(B,B);
    for(int l=0;l<B;l++) 
    {

        //AWL(i,j) = A[l,ij] * W[ij];
        Matrix ALW = W;
        for(int i=0;i<N;i++)
            for(int j=0;j<N;j++)
                if (not A(T,l,i,j))
                    ALW(i,j) = 0;

        vector<double> column = FastMTM(T,ALW,leaf_sets);

        for(int k=0;k<B;k++)
            M1(k,l) = column[k];
    }

    Matrix x = solve(M1,M2);

    assert(x.size1() == B);
    assert(x.size2() == 1);

    vector<double> b(B);
    for(int i=0;i<b.size();i++)
        b[i] = x(i,0);

    return b;
}

vector<double> LeastSquares(const Tree& T, const Matrix & D,
                            const vector<vector<int> >& leaf_sets) 
{
    const int N = T.n_leaves();

    Matrix I(N,N);
    for(int i=0;i<N;i++)
        for(int j=0;j<N;j++)
            I(i,j) = 1.0;

    return LeastSquares(T,D,I,leaf_sets);
}


vector<double> FastLeastSquares(const Tree& T, const Matrix & D,
                                const vector<vector<int> >& leaf_sets) 
{
    assert(is_Cayley(T));

    vector<double> delta = FastMTM(T,D,leaf_sets);

    vector<double> b(T.n_branches());

    int i=0;
    vector<const_branchview> branches;
    for(;i<T.n_leafbranches();i++) {
        // get the 2 neighboring branches, pointing *outwards*.
        branches.clear();
        append(T.directed_branch(i).branches_after(),branches);
        assert(branches.size() == 2);

        int j = branches[0].undirected_name();
        int k = branches[1].undirected_name();
        double Nj = leaf_sets[branches[0]].size();
        double Nk = leaf_sets[branches[1]].size();
    
        b[i] = (1+Nj+Nk)*delta[i]-(1+Nj-Nk)*delta[j] - (1-Nj+Nk)*delta[k];
        b[i] /= (4*Nj*Nk);
    }

    for(;i<T.n_branches();i++) {
        // get the 4 neighboring branches, pointing *outwards*.
        branches.clear();
        append(T.directed_branch(i).branches_after(),branches);
        append(T.directed_branch(i).branches_before(),branches);
        branches[2] = branches[2].reverse();
        branches[3] = branches[3].reverse();
        assert(branches.size() == 4);

        int j = branches[0].undirected_name();
        int k = branches[1].undirected_name();
        int l = branches[2].undirected_name();
        int m = branches[3].undirected_name();

        double Nj = leaf_sets[branches[0]].size();
        double Nk = leaf_sets[branches[1]].size();
        double Nl = leaf_sets[branches[2]].size();
        double Nm = leaf_sets[branches[3]].size();

        double N  = T.n_leaves();
        assert(Nj + Nk + Nl + Nm <= N);
    
        b[i] = (N/Nj + N/Nk + N/Nl + N/Nm - 4)*delta[i] +
            (Nj + Nk)/(Nj*Nk)*((2*Nk-N)*delta[j] + (2*Nj-N)*delta[k]) +
            (Nl + Nm)/(Nl*Nm)*((2*Nm-N)*delta[l] + (2*Nl-N)*delta[m]);
    
        b[i] /= (4*(Nj+Nk)*(Nl+Nm));
    }

    return b;
}


Matrix EdgesDistanceMatrix(const Tree& T) 
{
    const int N = T.n_nodes();
    const int n = T.n_leaves();
 
    vector<const_branchview> branches;
    branches.push_back(T.branch(0));
    branches.push_back(T.branch(0).reverse());

    list<int> points;

    Matrix D1(N,N);
    {
        int node = T.branch(0).target();
        D1(0,node) = D1(node,0) = 1;
    }
    
    while(not branches.empty()) 
    {
        int dead_point = branches.back().target();

        vector<const_branchview> new_branches;
        append(branches.back().branches_after(), new_branches);
        branches.pop_back();
    
        if (new_branches.empty()) {
            points.push_back(dead_point);
            continue;
        }

        for(int i=0;i<new_branches.size();i++) 
        {
            int p1 = new_branches[i].target();

            // Add distances from i -> leaves
            for(const auto& p2: points)
                D1(p1,p2) = D1(p2,p1) = D1(p2,dead_point) + 1;

            // Add distances from i -> active points 
            for(int j=0;j<branches.size();j++)
            {
                int p2 = branches[j].target();
                D1(p1,p2) = D1(p2,p1) = D1(p2,dead_point) + 1;
            }

            // Add distances from i -> new points
            for(int j=0;j<i;j++) 
            {
                int p2 = new_branches[j].target();
                D1(p1,p2) = D1(p2,p1) = 2;
            }
        }

        branches.insert(branches.end(), new_branches.begin(), new_branches.end());
    }

    assert(points.size() == T.n_leaves());
    for(const auto& p: points)
        assert(T.node(p).is_leaf_node());


    // Create the final matrix from the larger matrix
    Matrix D(n,n);
    for(int i=0;i<n;i++)
        D(i,i) = 0;

    for(int i=0;i<n;i++)
        for(int j=0;j<i;j++)
            D(i,j) = D(j,i) = D1(i,j);

    /*
      for(int i=0;i<n;i++)
      for(int j=0;j<n;j++) {
      if (not (std::abs(D(i,j) - T.edges_distance(i,j)) < 1.0e-8))
      std::cerr<<"D("<<i<<","<<j<<") = "<<D(i,j)<<"    T(i,j) = "<<T.edges_distance(i,j)<<endl;
      assert( std::abs(D(i,j) - T.edges_distance(i,j)) < 1.0e-8);
      }
    */

    return D;
}

Matrix DistanceMatrix(const Tree& T) 
{
    const int N = T.n_nodes();
    const int n = T.n_leaves();
 
    vector<const_branchview> branches;
    branches.push_back(T.branch(0));
    branches.push_back(T.branch(0).reverse());

    list<int> points;

    Matrix D1(N,N);
    {
        int node = T.branch(0).target();
        D1(0,node) = D1(node,0) = branches[0].length();
    }
    
    while(not branches.empty()) 
    {
        int dead_point = branches.back().target();

        vector<const_branchview> new_branches;
        append(branches.back().branches_after(), new_branches);
        branches.pop_back();
    
        if (new_branches.empty()) {
            points.push_back(dead_point);
            continue;
        }

        for(int i=0;i<new_branches.size();i++) 
        {
            int p1 = new_branches[i].target();

            double L = new_branches[i].length();

            // Add distances from i -> leaves
            for(const auto& p2: points)
                D1(p1,p2) = D1(p2,p1) = D1(p2,dead_point) + L;

            // Add distances from i -> active points 
            for(int j=0;j<branches.size();j++)
            {
                int p2 = branches[j].target();
                D1(p1,p2) = D1(p2,p1) = D1(p2,dead_point) + L;
            }

            // Add distances from i -> new points
            for(int j=0;j<i;j++) 
            {
                int p2 = new_branches[j].target();
                D1(p1,p2) = D1(p2,p1) = L + new_branches[j].length();
            }
        }

        branches.insert(branches.end(), new_branches.begin(), new_branches.end());
    }

    assert(points.size() == T.n_leaves());
    for(const auto& p: points)
        assert(T.node(p).is_leaf_node());


    // Create the final matrix from the larger matrix
    Matrix D(n,n);
    for(int i=0;i<n;i++)
        D(i,i) = 0;

    for(int i=0;i<n;i++)
        for(int j=0;j<i;j++)
            D(i,j) = D(j,i) = D1(i,j);

    /*
      for(int i=0;i<n;i++)
      for(int j=0;j<n;j++) {
      if (not (std::abs(D(i,j) - T.distance(i,j)) < 1.0e-8))
      std::cerr<<"D("<<i<<","<<j<<") = "<<D(i,j)<<"    T(i,j) = "<<T.distance(i,j)<<endl;
      assert( std::abs(D(i,j) - T.distance(i,j)) < 1.0e-8);
      }
    */
    return D;
}

Matrix C(const Matrix& M) {
    Matrix I = M;
    for(int i=0;i<I.size1();i++)
        for(int j=0;j<I.size2();j++)
            I(i,j) = 1.0 - I(i,j);

    return I;
}

double sum_1(const Matrix& D, int x)
{
    int N = D.size1();
    assert(N = D.size2());

    double total = 0;

    for(int i=0;i<N;i++)
        total += D(i,x);

    return total;
}

double sum_2(const Matrix& D, int x)
{
    int N = D.size1();
    assert(N = D.size2());

    double total = 0;

    for(int i=0;i<N;i++)
        total += D(x,i);

    return total;
}

std::pair<int,int> arg_min_dist(const Matrix& D)
{
    assert(D.size1() >= 2);
    assert(D.size1() == D.size2());

    int xmin = 0;
    int ymin = 1;
    double value = D(xmin,ymin);

    for(int i=0;i<D.size1();i++)
        for(int j=0;j<D.size2();j++)
            if (i != j and D(i,j) < value)
            {
                xmin = i;
                ymin = j;
                value = D(i,j);
            }

    return {xmin, ymin};
}

int cluster_edges(Tree& T, const tree_edge& e1, const tree_edge& e2)
{
    assert(e1.node1 != e2.node1);
    assert(e1.node2 == e2.node2);

    nodeview u = T.create_node_on_branch( T.branch(e1) );

    T.reconnect_branch(e2.node1, e2.node2, u);

    return u;
}


Tree NJ(Matrix D)
{
    int N = D.size1();

    // Handle N=0
    if (N == 0)
    {
        Tree T;
        return T;
    }

    // Handle N=2
    if (N == 1) {
        Tree T;
        T.add_first_node();
        return T;
    }

    // Handle N=2
    if (N == 2)
    {
        Tree T;
        T.add_first_node();
        T.add_leaf_node(0);
        T.branch(0).set_length(D(0,1));
        return T;
    }

    Tree T = star_tree(N);
    vector<tree_edge> mapping = T.leaf_branches();
    while (1)
    {
        N = D.size1();
        if (N == 3) break;

        vector<double> divergence(N);
        for(int i=0;i<N;i++)
        {
            divergence[i] = 0;
            for(int j=0;j<N;j++)
                if (i != j)
                    divergence[i] += D(i,j);
        }

        Matrix M(N,N);
        for(int i=0;i<N;i++)
            for(int j=0;j<N;j++)
            {
                M(i,j) = D(i,j) - (divergence[i] + divergence[j])/(N-2);
            }
    
        int f, g;
        tie(f,g) = arg_min_dist(M);
    
        //1. join the two leaves f and g to the new node u
        int u = cluster_edges(T, mapping[f], mapping[g]);
        tree_edge h(u, mapping[f].node2);
    
        //2. set the lengths of the edges from f->u and g->u
        T.directed_branch(mapping[f].node1, u).set_length( D(f,g) / 2 + (divergence[f]-divergence[g]) / (2*(N-2)) );
        T.directed_branch(mapping[g].node1, u).set_length( D(f,g) / 2 + (divergence[g]-divergence[f]) / (2*(N-2)) );
    
        //3. construct the new distance matrix and mapping 
    
        vector<int> new_to_old;
        vector<tree_edge> mapping2;
        for(int i=0;i<mapping.size();i++)
        {
            if (i == f or i ==g)
                ;
            else
            {
                new_to_old.push_back(i);
                mapping2.push_back(mapping[i]);
            }
        }
        new_to_old.push_back(-1);
        mapping2.push_back(h);
    
        assert(N-1 == new_to_old.size());
        Matrix D2(N-1, N-1);
    
        for(int i=0; i< N-1; i++)
            for(int j=0;j < N-1; j++)
            {
                int I = new_to_old[i];
                int J = new_to_old[j];
                if (I != -1)
                    assert(mapping2[i] == mapping[I]);
                if (J != -1)
                    assert(mapping2[j] == mapping[J]);
                if (I == -1 and J == -1)
                    D2(i,j) = 0;
                else if (I == -1)
                    D2(i,j) = (D(f,J) + D(g,J) - D(f,g))/2.0;
                else if (J == -1)
                    D2(i,j) = (D(I,f) + D(I,g) - D(f,g))/2.0;
                else
                    D2(i,j) = D(I,J);
            }

        D.resize(N-1, N-1);
        D = D2;
        mapping = mapping2;

        std::cout<<write_no_names(T)<<std::endl;
    }

    //5. if there is just one branch left, the set its length and quit.
    assert(D.size1() == 3);
    assert(mapping[0].node2 == mapping[1].node2);
    assert(mapping[0].node2 == mapping[2].node2);

    double L0 = ( D(0,1) + D(0,2) - D(1,2) ) / 2;
    double L1 = ( D(1,0) + D(1,2) - D(0,2) ) / 2;
    double L2 = ( D(2,0) + D(2,1) - D(0,1) ) / 2;

    T.directed_branch(mapping[0]).set_length( L0 );
    T.directed_branch(mapping[1]).set_length( L1 );
    T.directed_branch(mapping[2]).set_length( L2 );

    return T;
}
