#include <phaser/ncs/NCSSymmetry.h>
#include <phaser/io/Errors.h>

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/connected_components.hpp>

#include <scitbx/math/euler_angles.h>
#include <set>


namespace phaser {  namespace ncs {

// Utility function for set comparison
std::set< int >
set_intersect(const std::set< int >& left, const std::set< int >& right)
{
    std::set< int > intersect;

    for (
        std::set< int >::const_iterator it = left.begin();
        it != left.end();
        ++it
        )
    {
        if ( right.find( *it ) != right.end() )
        {
            intersect.insert( *it );
        }
    }

    return intersect;
}

//
// class NCSSymmetry
//

// Constructor
NCSSymmetry::NCSSymmetry(
    const std::vector< NCSOperator >& operators,
    double angular,
    double spatial
    )
    : angular_( angular ), spatial_( spatial )
{
    operators_.push_back( NCSOperator::unity );
    
    for ( opIter it = operators.begin(); it != operators.end(); ++it )
    {
        if ( !contains( *it ) )
        {
            operators_.push_back( *it );
        }
    }
}

// Destructor
NCSSymmetry::~NCSSymmetry()
{
}

// Accessors
const std::vector< NCSOperator >&
NCSSymmetry::get_operators() const
{
    return operators_;
}


std::vector< scitbx::vec3< double > >
NCSSymmetry::get_euler_list() const
{
    std::vector< scitbx::vec3< double > > euler_list;

    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        euler_list.push_back(
            scitbx::math::euler_angles::zyz_angles( it->get_rotation() )
            );
    }

    return euler_list;
}


std::vector< scitbx::vec3< double > >
NCSSymmetry::get_orth_list() const
{
    std::vector< scitbx::vec3< double > > orth_list;

    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        orth_list.push_back( it->get_translation() );
    }

    return orth_list;
}


double
NCSSymmetry::get_angular_tolerance() const
{
    return angular_;
}


double
NCSSymmetry::get_spatial_tolerance() const
{
    return spatial_;
}

// Display methods
std::string
NCSSymmetry::to_string() const
{
    std::ostringstream oss;

    oss << name() << " consisting of " << operators_.size() << " operators"
        << std::endl;

    for ( size_t index = 0; index < operators_.size(); ++index )
    {
        oss << std::endl << "Operator " << ( index + 1 ) << std::endl;
        oss << operators_[ index ].to_string() << std::endl;
    }

    return oss.str();
}


std::string
NCSSymmetry::name() const
{
    return "NCSSymmetry";
}

// Query methods
bool
NCSSymmetry::contains(const NCSOperator& op) const
{
    return find_operator( op ) != operators_.end();
}


bool
NCSSymmetry::equals(const std::vector< NCSOperator >& ops) const
{
    return ( operators_.size() == ops.size() )
        && ( operators_.size() == intersect( ops ).size() );
}

// Other methods
std::vector< NCSOperator >
NCSSymmetry::intersect(const std::vector< NCSOperator >& ops) const
{
    std::vector< NCSOperator > common_operators;

    for ( opIter it = ops.begin(); it != ops.end(); ++it )
    {
        if ( contains( *it ) )
        {
            common_operators.push_back( *it );
        }
    }

    return common_operators;
}


std::vector< NCSOperator >
NCSSymmetry::left_products(const NCSOperator& op) const
{
    std::vector< NCSOperator > products;
    
    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        products.push_back( op * ( *it ) );
    }

    return products;
}


std::vector< NCSOperator >
NCSSymmetry::right_products(const NCSOperator& op) const
{
    std::vector< NCSOperator > products;
    
    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        products.push_back( ( *it ) * op );
    }

    return products;
}


std::vector< std::pair< NCSSymmetry, NCSSymmetry > >
NCSSymmetry::factorize() const
{
    std::vector< std::set< int > > cross_closure_sets;
    std::vector< std::set< int > > all_closure_sets;

    for ( opIter oit = operators_.begin(); oit != operators_.end(); ++oit )
    {
        // Calculate operator set for which *oit is closed
        std::set< int > closure_relations;

        for ( opIter iit = operators_.begin(); iit != operators_.end(); ++iit )
        {
            // Closure wrt self is calculated later
            if ( oit == iit )
            {
                continue;
            }

            if ( contains( (*oit) * (*iit) ) )
            {
                closure_relations.insert( iit - operators_.begin() );
            }
        }

        //std::cout << oit->to_string();
        //std::cout << "Cross closure size: " << closure_relations.size()
        //    << std::endl;

        cross_closure_sets.push_back( closure_relations );

        if ( contains( (*oit) * (*oit) ) )
        {
            closure_relations.insert( oit - operators_.begin() );
        }

        //std::cout << "All closure size: " << closure_relations.size()
        //    << std::endl;
        all_closure_sets.push_back( closure_relations );
    }

    PHASER_ASSERT( cross_closure_sets.size() == all_closure_sets.size() );
    const int op_count = cross_closure_sets.size();

    // Create graph
    using namespace boost;
    adjacency_list< vecS, vecS, undirectedS > graph( op_count );

    // Add edge to graph if
    //   1. cross_closure_set(left).size() == cross_closure_set(right).size()
    //   2. all_closure_set(left) == all_closure_set(right)
    for ( int left = 0; left < op_count; ++left )
    {
        for ( int right = left+1; right < op_count; ++right )
        {
            const std::set< int >& l_cross = cross_closure_sets[ left ];
            const std::set< int >& r_cross = cross_closure_sets[ right ];
            const std::set< int >& l_all = all_closure_sets[ left ];
            const std::set< int >& r_all = all_closure_sets[ right ];
            std::set< int > common = set_intersect( l_all, r_all );

            if ( ( l_cross.size() == r_cross.size() )
                && ( common.size() == l_all.size() )
                && ( common.size() == r_all.size() ) )
            {
                add_edge( left, right, graph );
            }
        }
    }

    // Find connected components
    std::vector< int > component_assignments( op_count );
    int comp_count = connected_components( graph, &component_assignments[0] );

    //std::cout << comp_count << " components found" << std::endl;

    // Collect operators for each component, determine factors and store unique
    std::vector< std::pair< NCSSymmetry, NCSSymmetry > > factorizations;

    for ( int comp_index = 0; comp_index < comp_count; ++comp_index )
    {
        std::vector< NCSOperator > candidates;
        std::vector< std::set< int > > cc_sets;

        for ( int op_index = 0; op_index < op_count; ++op_index )
        {
            if ( component_assignments[ op_index ] == comp_index )
            {
                candidates.push_back( operators_[ op_index ] );
                cc_sets.push_back( cross_closure_sets[ op_index ] );
            }
        }

        PHASER_ASSERT( 0 < cc_sets.size() );

        // Calculate intersect of cross_closure_sets
        std::set< int > cc_intersection = cc_sets[0];

        for (
            std::vector< std::set< int > >::const_iterator it =
                cc_sets.begin() + 1;
            it != cc_sets.end();
            ++it
            )
        {
            cc_intersection = set_intersect( cc_intersection, *it );
        }

        // Collect the operators
        std::vector< NCSOperator > cc_int_ops;

        for(
            std::set<int>::const_iterator it = cc_intersection.begin();
            it != cc_intersection.end();
            ++it
            )
        {
            PHASER_ASSERT( *it < op_count );
            cc_int_ops.push_back( operators_[ *it ] );
        }
        
        // Select outer symmetry operators by using cc_intersection as an
        // equivalence class
        NCSSymmetry equiv_class( cc_int_ops, angular_, spatial_ );
        std::vector< NCSOperator > accounteds;
        std::vector< NCSOperator > outer_factor_ops;

        for ( opIter it = candidates.begin(); it != candidates.end(); ++it )
        {
            NCSSymmetry sym( accounteds, angular_, spatial_ );

            if ( sym.contains( *it ) )
            {
                continue;
            }

            outer_factor_ops.push_back( *it );
            accounteds.push_back( *it );

            std::vector< NCSOperator > products =
                equiv_class.left_products( *it );
            accounteds.insert(
                accounteds.end(),
                products.begin(),
                products.end()
                );
        }
        
        // Calculate outer and inner
        NCSSymmetry outer = NCSSymmetry( outer_factor_ops, angular_, spatial_ );
        NCSSymmetry inner = NCSSymmetry( quotient( outer ), angular_, spatial_ );

        // Check whether factors produce the symmetry
        std::vector< NCSOperator > dirprod = outer.direct_product( inner );
        std::pair< NCSSymmetry, NCSSymmetry > factors = 
            std::pair< NCSSymmetry, NCSSymmetry >(
                NCSSymmetry( std::vector< NCSOperator >(), angular_, spatial_ ),
                *this
                );

        if ( ( operators_.size() == dirprod.size() )
            && ( operators_.size() == intersect( dirprod ).size() ) )
        {
            factors = std::pair< NCSSymmetry, NCSSymmetry >( outer, inner );
        }

        // Store factors if unique
        bool unique = true;

        for(
            std::vector< std::pair< NCSSymmetry, NCSSymmetry > >
                ::const_iterator it = factorizations.begin();
            it != factorizations.end();
            ++it
            )
        {
            if ( factors.first.equals( it->first.get_operators() ) )
            {
                unique = false;
                break;
            }
        }

        if ( unique )
        {
            factorizations.push_back( factors );
        }
    }

    return factorizations;
}


std::vector< NCSOperator >
NCSSymmetry::quotient(const NCSSymmetry& divisor) const
{
    std::vector< NCSSymmetry > equivs;
    std::vector< NCSOperator > ops;

    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        std::vector< NCSOperator > prods = divisor.right_products( *it );
        
        if ( intersect( prods ).size() != prods.size() )
        {
            continue;
        }

        bool has_class = false;

        for (
            std::vector< NCSSymmetry >::const_iterator eqit = equivs.begin();
            eqit != equivs.end();
            ++eqit
            )
        {
            if ( eqit->intersect( prods ).size() == prods.size() )
            {
                has_class = true;
                break;
            }
        }

        if ( !has_class )
        {
            equivs.push_back( NCSSymmetry( prods, angular_, spatial_ ) );
            ops.push_back( *it );
        }
    }

    return ops;
}


std::vector< NCSOperator >
NCSSymmetry::direct_product(const NCSSymmetry& inner) const
{
    std::vector< NCSOperator > ops;

    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        std::vector< NCSOperator > prods = inner.left_products( *it );
        ops.insert( ops.end(), prods.begin(), prods.end() );
    }

    return ops;
}

// Protected methods
NCSSymmetry::opIter
NCSSymmetry::find_operator(const NCSOperator& op) const
{
    NCSOperator inverse( op.inverse() );

    for ( opIter it = operators_.begin(); it != operators_.end(); ++it )
    {
        if ( inverse.approx_inverse( *it, angular_, spatial_ ) )
        {
            return it;
        }
    }

    return operators_.end();
}

}  } // namespace phaser::ncs