/*
 * Copyright (c) 1999 Johnny C. Lam.
 * All rights reserved.
 *
 * Permission to use, copy, modify, distribute and sell this software
 * and its documentation for any purpose is hereby granted without fee,
 * provided that the above copyright notice appear in all copies and
 * that both that copyright notice and this permission notice appear
 * in supporting documentation.  The author makes no representations
 * about the suitability of this software for any purpose.  It is
 * provided "as is" without express or implied warranty.
 */

#include "latent_recomb_strategy.h"

#include <algorithm>
#include <deque>
#include <limits>
#include <memory>
#include <vector>

#include "haplotype.h"
#include "probability.h"
#include "types.h"

class HaplotypeChange {
public:
    typedef Haplotype::size_type size_type;

public:
    HaplotypeChange(size_type i) : marker_index(i) {}

    size_type marker_index;
    virtual void Change(Haplotype&, const Haplotype&) const = 0;
};

class Mutation : public HaplotypeChange {
public:
    Mutation(size_type i) : HaplotypeChange(i) {}

    virtual void Change(Haplotype& s, const Haplotype& t) const {
	s.Mutate(marker_index, t[marker_index]);
    }
};

class LeftRecombination : public HaplotypeChange {
public:
    LeftRecombination(size_type i) : HaplotypeChange(i) {}

    virtual void Change(Haplotype& s, const Haplotype& t) const {
	s.LeftRecombine(marker_index + 1, t);
    }
};

class RightRecombination : public HaplotypeChange {
public:
    RightRecombination(size_type i) : HaplotypeChange(i) {}

    virtual void Change(Haplotype& s, const Haplotype& t) const {
	s.RightRecombine(marker_index, t);
    }
};

/*
 * Assumes [first_mut, last_mut) is sorted range.
 */

template <class RandomAccessIterator>
Path
GenPath(const Haplotype &s, const Haplotype &t,
	PathStrategy::size_type disease,
	RandomAccessIterator first_change,
	RandomAccessIterator last_change,
	Probability *rules, bool verbose)
{
    std::vector<const HaplotypeChange*> ch(last_change - first_change);
    std::vector<const HaplotypeChange*>::iterator copy_it = ch.begin();
    while (first_change != last_change) {
	*copy_it = *first_change;
	++copy_it;
	++first_change;
    }
    
    /* Initialize min_path with solely recomb from s to t. */
    Path min_path(s);
    min_path.Add(Path::Edge(t, rules->log_P(s, t, disease)));
    unsigned int min_num_latents = 0;
    log_probability min_path_weight = min_path.Weight();

    /* Find the path of minimum weight from s to t. */
    do {
	if (verbose) {
	    std::cout << "    " << "Index permutation:";
	    std::vector<const HaplotypeChange*>::const_iterator f = ch.begin();
	    std::vector<const HaplotypeChange*>::const_iterator end = ch.end();
	    while (f != end) {
		std::cout << " " << (*f)->marker_index;
		++f;
	    }
	    std::cout << endl;
	}

	Path path(s);	// Path begins with s
	std::vector<const HaplotypeChange*>::const_iterator f = ch.begin();
	Haplotype prev(s);
	unsigned int num_latents = 0;
	log_probability path_weight = 0;

	while (f != ch.end()) {
	    Haplotype latent(prev);
	    (*f)->Change(latent, t);
	    Haplotype::property_t properties = latent.Properties();
	    if (!(properties & Haplotype::DISEASED))
		latent.SetProperties(properties | Haplotype::LATENT);
	    log_probability latent_prob = rules->log_P(prev, latent, disease);
	    path.Add(Path::Edge(latent, latent_prob));
	    if (latent.IsLatent())
		++num_latents;
	    path_weight += latent_prob;
	    if (path_weight >= min_path_weight)
		break;
	    prev = latent;
	    ++f;
	}
	if (verbose &&
	    path_weight != std::numeric_limits<log_probability>::max()) {
	    std::cout << "\t" << "Test path:" << endl;
	    std::cout << "\t" << "        " << s << endl;
	    Path::const_iterator pf = path.begin(), pe = path.end();
	    ++pf;
	    while (pf != pe) {
		std::cout << "\t" << "  ----> " << pf->haplotype
			  << " :\t" << pf->weight << endl;
		++pf;
	    }
	    std::cout << "\t" << "Test path weight (-log(prob)): "
		      << path_weight << endl;
	}
	if ((path_weight < min_path_weight) ||
	    ((path_weight == min_path_weight) &&
	     (num_latents < min_num_latents))) {
	    min_path = path;
	    min_num_latents = num_latents;
	    min_path_weight = path_weight;
	    if (verbose) {
		std::cout << "\t" << "New min path weight: "
			  << min_path_weight << endl;
	    }
	}
    } while (next_permutation(ch.begin(), ch.end()));

    return min_path;
}

Path
LatentRecombStrategy::MakePath(const Haplotype &s, const Haplotype &t,
			       size_type disease, Probability* rules) const
{
    if (Verbose()) {
	std::cout << "\t" << "Source: " << s << endl
		  << "\t" << "Target: " << t << endl
		  << "\t" << "Disease: " << disease << endl;
    }

    typedef std::vector<size_type> IndexList;
    IndexList left_indices, right_indices;
    std::deque<size_type> indices;

    /* IndexList needs to have bidirectional iterators */
    unsigned int left_mask = rules->LeftMask();
    unsigned int right_mask = rules->RightMask();

    /* Make a list of the indices where s,t differ. */
    for (size_type i = disease; i != left_mask; --i) {
	if (s[i - 1] != t[i - 1]) {
	    left_indices.push_back(i - 1);
	    indices.push_front(i - 1);
	}
    }
    for (size_type i = disease; i < right_mask; ++i) {
	if (s[i] != t[i]) {
	    right_indices.push_back(i);
	    indices.push_back(i);
	}
    }
    if (Verbose()) {
	std::cout << "    "  << "Left indices:";
	for (unsigned int i = 0; i < left_indices.size(); ++i)
	    std::cout << " " << left_indices[i];
	std::cout << endl;
	std::cout << "    "  << "Right indices:";
	for (unsigned int i = 0; i < right_indices.size(); ++i)
	    std::cout << " " << right_indices[i];
	std::cout << endl;
	std::cout << "    "  << "Indices:";
	for (unsigned int i = 0; i < indices.size(); ++i)
	    std::cout << " " << indices[i];
	std::cout << endl;
    }

    Path min_path;
    log_probability min_path_weight =
	std::numeric_limits<log_probability>::max();

    if (left_indices.empty()) {
	if (Verbose())
	    std::cout << "    " << "Right changes only." << endl;
	unsigned int i = 0;
	for (unsigned int j = 0; j < right_indices.size(); ++j) {
	    if (j - i + 1 < _max_latents + 1) {
		vector<HaplotypeChange*> changes(j - i + 1);
		for (unsigned int k = i; k < j; ++k)
		    changes[k - i] = new Mutation(right_indices[k]);
		changes[j - i] = new RightRecombination(right_indices[j]);
		Path path = GenPath(s, t, disease,
				    changes.begin(), changes.end(),
				    rules, Verbose());
		log_probability path_weight = path.Weight();
		if ((path_weight < min_path_weight) ||
		    (min_path_weight
		     == std::numeric_limits<log_probability>::max())) {
		    min_path = path;
		    min_path_weight = path_weight;
		for (unsigned int k = i; k <= j; ++k)
		    delete changes[k - i];
		}
	    }
	}
    } else if (right_indices.empty()) {
	if (Verbose())
	    std::cout << "    " << "Left changes only." << endl;
	unsigned int i = 0;
	for (unsigned int j = 0; j < left_indices.size(); ++j) {
	    if (j - i + 1 < _max_latents + 1) {
		std::vector<HaplotypeChange*> changes(j - i + 1);
		for (unsigned int k = i; k < j; ++k)
		    changes[k - i] = new Mutation(left_indices[k]);
		changes[j - i] = new LeftRecombination(left_indices[j]);
		Path path = GenPath(s, t, disease,
				    changes.begin(), changes.end(),
				    rules, Verbose());
		log_probability path_weight = path.Weight();
		if ((path_weight < min_path_weight) ||
		    (min_path_weight
		     == std::numeric_limits<log_probability>::max())) {
		    min_path = path;
		    min_path_weight = path_weight;
		for (unsigned int k = i; k <= j; ++k)
		    delete changes[k - i];
		}
	    }
	}
    } else {
	if (Verbose())
	    std::cout << "    " << "Left and right changes." << endl;

	// Index into indices where disease should be.
	unsigned int d_index = left_indices.size();

	for (unsigned int i = 0; i < d_index; ++i) {
	    for (unsigned int j = d_index; j < indices.size(); ++j) {
		if (j - i + 1 < _max_latents + 1) {
		    std::vector<HaplotypeChange*> changes(j - i + 1);
		    changes[i - i] = new LeftRecombination(indices[i]);
		    for (unsigned int k = i + 1; k < j; ++k)
			changes[k - i] = new Mutation(indices[k]);
		    changes[j - i] = new RightRecombination(indices[j]);
		    Path path = GenPath(s, t, disease,
					changes.begin(), changes.end(),
					rules, Verbose());
		    log_probability path_weight = path.Weight();
		    if ((path_weight < min_path_weight) ||
			(min_path_weight
			 == std::numeric_limits<log_probability>::max())) {
			min_path = path;
			min_path_weight = path_weight;
		    for (unsigned int k = i; k <= j; ++k)
			delete changes[k - i];
		    }
		}
	    }
	}
    }
    return min_path;
}
