#pragma once

#include "dominator_list.hpp"
#include "address_list.hpp"
#include "chocolat/string/string_repeat.hpp"

#include <vector>
#include <map>
#include <boost/shared_ptr.hpp>

class AddressArrowList
{
public:
	typedef std::multimap<unsigned, unsigned> multimap;
	typedef std::pair<unsigned, unsigned> pair;
	typedef multimap::iterator iterator;
	typedef boost::shared_ptr<AddressArrowList> shared_ptr;
	
	class range_iterator
	{
		multimap::iterator m_iterator;
	public:
		range_iterator(multimap::iterator i)
				:m_iterator(i)
		{
		}
		range_iterator(const range_iterator& that)
				:m_iterator(that.m_iterator)
		{
		}
		unsigned operator*()const
		{
			return m_iterator->second;
		}
		range_iterator& operator++()
		{
			++m_iterator;
			return *this;
		}
		range_iterator& operator=(const range_iterator& rh)
		{
			m_iterator = rh.m_iterator;
			return *this;
		}
		bool operator==(const range_iterator& rh)const
		{
			return m_iterator == rh.m_iterator;
		}
		bool operator!=(const range_iterator& rh)const
		{
			return !(*this == rh);
		}
	};
	/** treat map equal range as list.*/
	class range
	{
		range_iterator m_begin;
		range_iterator m_end;
		unsigned m_size;
	public:
		range(multimap::iterator begin, multimap::iterator end)
				:m_begin(begin),
				 m_end(end)
		{
			m_size = 0;
			for (auto i=begin; i!=end; ++i) {
				++m_size;
			}
		}
		unsigned size()const
		{
			return m_size;
		}
		bool empty()const
		{
			return m_size == 0;
		}
		range_iterator begin()
		{
			return m_begin;
		}
		range_iterator end()
		{
			return m_end;
		}
	};
	multimap m_arrow;
	multimap m_inverse;
	static shared_ptr create()
	{
		shared_ptr result(new AddressArrowList);
		return result;
	}
	void insertArrow(unsigned head, unsigned tail) 
	{
		m_arrow.insert(pair(head,tail));
		m_inverse.insert(pair(tail,head));
	}
	range tails_of(unsigned address)
	{
		auto iterator_pair = m_arrow.equal_range(address);
		range result(iterator_pair.first, iterator_pair.second);
		return result;
	}
	range heads_of(unsigned address)
	{
		auto iterator_pair = m_inverse.equal_range(address);
		range result(iterator_pair.first, iterator_pair.second);
		return result;
	}
};
/** */
class SentenceStructure
{
public:
	typedef boost::shared_ptr<SentenceStructure> shared_ptr;
	unsigned m_header;
	BasicBlockList::shared_ptr m_block_list;
	
	SentenceStructure(BasicBlock::shared_ptr header)
			:m_header(header->startAddress()),
			 m_block_list(BasicBlockList::create())
	{
	}
	virtual ~SentenceStructure()
	{
	}
	virtual void display(std::ostream& os)const
	{
	}
	void push(BasicBlock::shared_ptr block)
	{
		(*m_block_list)[block->startAddress()] = block;
	}
	bool find(BasicBlock::shared_ptr block)const
	{
		return m_block_list->find(block->startAddress())
				!= m_block_list->end();
	}
};
class BranchStructure : public SentenceStructure
{
	BranchStructure(BasicBlock::shared_ptr header)
			:SentenceStructure(header)
	{
	}
	void construct()
	{
	}
};
/** */
class LoopStructure : public SentenceStructure
{
public:
	typedef boost::shared_ptr<LoopStructure> shared_ptr;
	
	/** predessor -> entry */
	AddressMarker::shared_ptr m_entry_list;
	/** break -> successor */
	AddressMarker::shared_ptr m_break_list;
	/** header <- continue */
	AddressMarker::shared_ptr m_continue_list;
	
	AddressMarker::shared_ptr m_successor_list;
	AddressMarker::shared_ptr m_predessor_list;
	
	
	LoopStructure(BasicBlock::shared_ptr header)
			:SentenceStructure(header)
	{
		m_entry_list.reset(new AddressMarker);
		m_break_list.reset(new AddressMarker);
		m_continue_list.reset(new AddressMarker);
		m_successor_list.reset(new AddressMarker);
		m_predessor_list.reset(new AddressMarker);
		
		push(header);
	}
	virtual void display(std::ostream& os)const
	{
		displayBriefResult(os);
		const auto npos = m_block_list->end();
		for (auto i=m_block_list->begin(); i != npos; ++i) {
			unsigned addr = i->first;
			if (addr == m_header) {
				std::cout << "DO    ";
			} else if (m_entry_list->found(addr)) {
				std::cout << "  ENT ";
			} else {
				std::cout << "      ";
			}
			std::cout << HexadecimalFormat(addr) << " ";
			bool cnt = m_continue_list->found(addr);
			bool brk = m_break_list->found(addr);
			
			if (cnt && brk) {
				std::cout << "WHILE";
			} else if (cnt) {
				std::cout << "CONTINUE";
			} else if (brk) {
				std::cout << "BREAK";
			}
			std::cout << std::endl;
		}
	}
	void displayBriefResult(std::ostream& os)const
	{
		os << HexadecimalFormat(m_header) << ": ";
		os << "pre " << m_predessor_list->size() << ", ";
		os << "ent " << m_entry_list->size() << ", ";
		os << "cnt " << m_continue_list->size() << ", ";
		os << "brk " << m_break_list->size() << ", ";
		os << "suc " << m_successor_list->size() << ": ";
		
		for (auto i=m_block_list->begin(); i != m_block_list->end(); ++i) {
			os << HexadecimalFormat(i->first) << " ";
		}
		os << std::endl;
	}
	/** search the continue and break blocks after all block are pushed. */
	void construct()
	{
		const auto npos = m_block_list->end();
		for (auto i=m_block_list->begin(); i != m_block_list->end(); ++i) {
			auto srcs = i->second->predessorList();
			for (auto j=srcs->begin(); j != srcs->end(); ++j) {
				if (m_block_list->find(*j) == npos){
					m_entry_list->mark(i->first);
					m_predessor_list->mark(*j);
				} 
			}
			auto dests = i->second->successorList();
			for (auto j=dests->begin(); j != dests->end(); ++j) {
				if (*j == m_header) {
					m_continue_list->mark(*j);
				} else if (m_block_list->find(*j) == npos){
					m_break_list->mark(i->first);
					m_successor_list->mark(*j);
				}
			}
		}
		// TODO: sort by depth first order.
	}
};
class LoopSet
{
public:
	std::vector<SentenceStructure::shared_ptr> m_vector;
	BasicBlockList::shared_ptr m_block_list;
	DominatorLattice::shared_ptr m_dominators;
	
	LoopSet(BasicBlockList::shared_ptr block_list,
			DominatorLattice::shared_ptr dominators)
			:m_block_list(block_list),
			 m_dominators(dominators)
	{
	}
	/** natural loop has only one header block.*/
	void addNaturalLoop(BasicBlock::shared_ptr header,
						BasicBlock::shared_ptr tail) 
	{
		BasicBlockBox work_list;
		LoopStructure::shared_ptr loop(new LoopStructure(header));
		m_vector.push_back(loop);
		
		if (header->startAddress() != tail->startAddress()) {
			loop->push(tail);
			work_list.push(tail);
			
			while (!work_list.empty()){
				BasicBlock::shared_ptr block = work_list.pop();
				AddressList::shared_ptr list = block->predessorList();
				for (auto i=list->begin(); i != list->end(); ++i) {
					BasicBlock::shared_ptr pred = (*m_block_list)[*i];
					if (!loop->find(pred)){
						loop->push(pred);
						work_list.push(pred);
					}
				}
			}
		}
		loop->construct();
	}
	void display(std::ostream& os)const
	{
		os << m_vector.size() << " natural loops found." << std::endl;
		for (unsigned i=0; i<m_vector.size(); ++i) {
			m_vector[i]->display(os);
			os << std::endl;
		}
	}
};
/** this worker find loop structure in the routine.*/
class LoopAnalyzer
{
	Routine::shared_ptr m_routine;
	DominatorLattice::shared_ptr m_dominators;
	AddressMarker m_visits;
	unsigned m_current_depth;
public:
	LoopAnalyzer()
	{
	}
	void initialize(Routine::shared_ptr routine,
					DominatorLattice::shared_ptr dominators)
	{
		m_routine = routine;
		m_dominators = dominators;
		m_visits.clear();
		m_current_depth = 0;
		
	}
	void analyze(Routine::shared_ptr routine,
				 DominatorLattice::shared_ptr dominators)
	{
		initialize(routine, dominators);
		//analyzeDepthFirstOrder(m_routine->blockList()->begin()->first);
		//return;

		
		BasicBlockList::shared_ptr block_list = routine->blockList();
		LoopSet loop_set(block_list, dominators);
		AddressListManipulater entries(routine->entryPointList());
		
		for (auto i=block_list->begin(); i!=block_list->end(); ++i) {
			BasicBlock::shared_ptr block = i->second;
			/*
			std::cout << std::endl;
			std::cout << HexadecimalFormat(i->first) << ":";
			//*/
			for (auto succ=block->successorList()->begin();
				 succ != block->successorList()->end(); ++succ)
			{
				/*
				std::cout << " " << HexadecimalFormat(*succ);
				//*/
				if (dominators->of(i->first)->find(*succ) ) {
					/*
					std::cout << "!";
					//*/
					BasicBlock::shared_ptr successor = (*block_list)[*succ];
					loop_set.addNaturalLoop(successor, block);
				}
			}
		}
		// loop_set.display(std::cout);
	}
	/* it is not used now.*/
	void analyzeDepthFirstOrder(unsigned start_address)
	{
		BasicBlockList list(m_routine->blockList());
		BasicBlock::shared_ptr block=list.get(start_address);
		if (block->isValid() && !m_visits.found(start_address)) {
			m_visits.mark(block->startAddress());
			++m_current_depth;
			std::cout << DecimalFormat(m_current_depth).width(3) << ": "
					  << HexadecimalFormat(start_address) << std::endl;
			
			auto successorList = block->successorList();
			if (successorList->size()==2){
				for (auto i =successorList->begin();
					 i != successorList->end(); ++i) {
					analyzeDepthFirstOrder(*i);
				}
			} else if (successorList->size()==1) {
				analyzeDepthFirstOrder(*successorList->begin());
			}
			--m_current_depth;
		} 
	}
};

