//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program 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 Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program.  If not, see http://www.gnu.org/licenses/.
//

#include "LEACH.h"
#include "FindModule.h"
#include "Distance.h"
#include <cmath>
#include <limits>
#include <LEACH_TDMA_Schedule_m.h>
#include <BatteryCSMAMacLayer.h>

Define_Module(LEACH);

void LEACH::initialize(int stage) {
	BaseLayer::initialize(stage);

	if(stage==0){
		headerLength= par("headerLength");
		arp = BaseArpAccess().get();
		myNetwAddr = this->findHost()->par("isBS") ?  L3BS : arp->myNetwAddr(this);

		EV << " myNetwAddr " << myNetwAddr << endl;

		lastRoundCH = -1;
		currentRound = -1;
		droppedMsgs = 0;

		P = par("P"); //percentage of CHs
		roundTime = par("roundTime");
		slotTime = par("slotTime");
		compressionIndex = par("compressionIndex");
		waitingTime = par("watingTime");
		MAX_CLUSTER_SIZE = par("maxClusterSize");
		timerNextStage = NULL;
		packetToSend = NULL;
	}

	if(stage==1)
	{
		//Round will start in 'roundTime' seconds
		timerNextRound = new cMessage("New Round", TIMER_NEW_ROUND);
		scheduleAt(simTime() + roundTime, timerNextRound);
	}
}

/**
 * Decapsulates the packet from the received Network packet
 **/
void LEACH::setHighTxPower()
{
	BatteryCSMAMacLayer* layer = (BatteryCSMAMacLayer*)(this->findHost()->getSubmodule("nic")->getSubmodule("mac"));
	layer->setToHighTwPower();
}

void LEACH::setLowTxPower()
{
	BatteryCSMAMacLayer* layer = (BatteryCSMAMacLayer*)(this->findHost()->getSubmodule("nic")->getSubmodule("mac"));
	layer->setToLowTxPower();
}

double LEACH::calculateThreshold()
{
	double threshold = 0;

	if((lastRoundCH == -1) || (currentRound - lastRoundCH > 1/P));
		threshold = P/(1-P*fmod(currentRound, 1/P));

	return threshold;
}

//--------------------------------------------------


//Stage for all the nodes: election of CHs

void LEACH::advertisementPhase()
{
	currentPhase = ADVERTISEMENT;

	//Code for better visualization into Tkenv environment
	this->findHost()->setDisplayString("i=misc/node,green;is=vs");
	this->findHost()->bubble("ADVERTISEMENT PHASE");

	//TODO: Clear phases
	distanceToCHs.clear();
	membersCH.clear();
	netwQueue.clear();

	if(dblrand() < calculateThreshold())
	{
		//I'm Cluster Head, I must broadcast my CH status,
		//ignore CH announces and WAIT for CH join messages.
		//"My CH" will be the BS
		lastRoundCH = currentRound;
		myCH = L3BS;
		this->findHost()->setDisplayString("i=misc/node,blue;is=vs");

		//Broadcast my CH status
	    NetwPkt* pkt = new NetwPkt("I'm Cluster Head", CH_STATUS_BROADCAST);
	    pkt->setDestAddr(L3BROADCAST);
	    pkt->setSrcAddr(myNetwAddr);
	    pkt->setBitLength(headerLength);
	    pkt->setControlInfo( new MacControlInfo(L2BROADCAST) );

	    EV << "Sending CH status broadcast\n";
	    sendDown(pkt);

	    //Wait for Join CH messages
	    currentPhase = CLUSTER_SETUP;
	    this->findHost()->bubble("CLUSTER_SETUP PHASE");
	}

	else
	{	//Non-CH: Waits to receive CH announces
		timerNextStage = new cMessage("Timer for join to a CH", TIMER_JOIN_CH);
		scheduleAt(simTime() + waitingTime, timerNextStage);
	}
}


//Stage for Non-CHs: send 'Join CH' message to the elected CH

void LEACH::clusterSetUpPhase()
{
	//Sends "join CH" message and waits for the TDMA schedule packet
	currentPhase = CLUSTER_SETUP;
	this->findHost()->bubble("CLUSTER_SETUP PHASE");

	DistanceToCHsMap::const_iterator it = distanceToCHs.begin();
	double minDistance = std::numeric_limits<double>::max();
	int nodeMinDistanceAddr = -1;

	//Get the minimun distance CH
	for(; it != distanceToCHs.end(); ++it)
	{
		if (it->second < minDistance)
		{
			minDistance = it->second;
			nodeMinDistanceAddr = it->first;
		}
	}

	//Clear 'the distanceToCHs' map
	distanceToCHs.clear();

	//Send join message and wait for TDMA schedule
	myCH = nodeMinDistanceAddr;

	NetwPkt* pkt = new NetwPkt("I want to join to your Cluster", JOIN_CH);
	pkt->setDestAddr(myCH);
	pkt->setSrcAddr(myNetwAddr);
	pkt->setBitLength(headerLength);
	//int macAddr = arp->getMacAddr(myCH);
	pkt->setControlInfo( new MacControlInfo(arp->getMacAddr(myCH)) );

	EV << "Sending join CH status to node[" << simulation.getModule(myCH)->getParentModule()->getIndex() << "]\n";
	sendDown(pkt);

	//Waits for the TDMA schedule
	currentPhase = SCHEDULE_CREATION;
	this->findHost()->bubble("SCHEDULE_CREATION PHASE");
}


//Stage for CHs: create the TDMA schedule and broadcast it

void LEACH::scheduleCreation()
{
	//Creates the message with the TDMA schedule and broadcasts it
	//After it, receives all the messages from the child nodes and stores them

	currentPhase = SCHEDULE_CREATION;
	this->findHost()->bubble("SCHEDULE_CREATION PHASE");

    LEACH_TDMA_Schedule* sched_pkt = new LEACH_TDMA_Schedule("TDMA Schedule", TDMA_SCHEDULE);
    sched_pkt->setDestAddr(L3BROADCAST);
    sched_pkt->setSrcAddr(myNetwAddr);
    //The length is supposed to be the header length + data length: num. cluster members * netw_address_size
    sched_pkt->setBitLength(headerLength + membersCH.size()*12);
    sched_pkt->setControlInfo( new MacControlInfo(L2BROADCAST) );
    sched_pkt->setScheduleArraySize(membersCH.size());

    MembersCHList::const_iterator it = membersCH.begin();

    for(int i=0; it!=membersCH.end(); ++it){
    	sched_pkt->setSchedule(i, *it);
    	i++;
    }

    EV << "Sending TDMA Schedule broadcast\n";
    sendDown(sched_pkt);

    //Wait until all Non-CH have sent data to me
    //TODO: DATA_TRANSMISSION PHASE ELIMINATED
//    timerNextRound = new cMessage("Timer to wait for CH-BS transfer", TIMER_SEND_DATA);
//    scheduleAt(simTime() + MAX_CLUSTER_SIZE* slotTime, timerNextRound);
//    currentPhase = DATA_TRANSMISSION;
//    this->findHost()->bubble("DATA_TRANSMISSION PHASE");
}


//Stage for all nodes: 1st Non-CHs send data to CH and after it CHs send data to BS

void LEACH::dataTransmission()
{
	//If I'm Non-CH
	if(lastRoundCH != currentRound){
		currentPhase = DATA_TRANSMISSION;
		this->findHost()->bubble("DATA_TRANSMISSION PHASE");
		simtime_t before_sending_time = simTime();

		//Sends its packet to the CH
		//After sending the messages, do nothing until next round
		packetToSend->setDestAddr(myCH);
		packetToSend->setControlInfo( new MacControlInfo(arp->getMacAddr(myCH)) );
		sendDown(packetToSend);
		packetToSend = NULL;

		//Checks if it works properly
		if(simTime() > before_sending_time + slotTime) {
			EV << "=============> Error!!! Time spent for sending message higher than expected" << endl;
			opp_error("ERROR: Time spent for sending message higher than expected");
		}
	}

	//If I'm CH: send data compressed to the BS:
	//Simulation of data compression
	else{
		int dataLength = 0;
		NetwQueueList::const_iterator it = netwQueue.begin();

		for( ; it!=netwQueue.end(); ++it){
			dataLength += (*it)->getBitLength();
		}

		if(packetToSend != NULL)
			dataLength += packetToSend->getBitLength();

		NetwPkt *pkt = new NetwPkt("Compressed data from the whole Cluster", (*it)->getKind());
		pkt->setBitLength((int)(ceil(dataLength*(1-compressionIndex))));
		pkt->setSrcAddr(myNetwAddr);
		pkt->setDestAddr(L3BS);
		pkt->setControlInfo(new MacControlInfo(L2BS));

		//Changes to High TX Power, sends the packet and returns to Low TX Power
		this->setHighTxPower();
		sendDown(pkt);
		this->setLowTxPower();

		//Cleans NetwQueue and packetToSend
		netwQueue.clear();
		delete packetToSend;
		packetToSend = NULL;

		//Waits until the next round
		this->findHost()->setDisplayString("i=misc/node,green;is=vs");
	}
}













void LEACH::handleLowerMsg(cMessage* msg){

	switch( msg->getKind() ){

	case CH_STATUS_BROADCAST:{
		//If I'm Non-CH, I save the announce
		//If not, I ignore the message

		if(lastRoundCH != currentRound){
			if(currentPhase != ADVERTISEMENT)
			{	//If we receive this message in a different phase -> Ignore it
				EV << "ERROR: 'CH Status Broadcast' received during phase " << currentPhase << endl;
			}

			else
			{	//Save the CH candidate if it's in range
				NetwPkt *m = static_cast<NetwPkt *>(msg);
				double distance = Distance::getDistanceByNetwAddr(myNetwAddr, m->getSrcAddr());
				distanceToCHs[m->getSrcAddr()] = distance;

				//Reactivate the timer and wait for other announce or the next phase
				cancelAndDelete(timerNextStage);
				timerNextStage = new cMessage("Timer for join to a CH", TIMER_JOIN_CH);
				scheduleAt(simTime() + waitingTime, timerNextStage);
			}
		}
		delete msg;
	}
	break;

	case JOIN_CH:	{

		if(currentPhase != CLUSTER_SETUP)
		{	//If we receive this message in a different phase -> Ignore it
			EV << "ERROR: 'Join CH' received during phase " << currentPhase << endl;
		}

		else {
			//Of course I'm CH if MAC layer gives me the message
			//I must save the node member address
			NetwPkt *m = static_cast<NetwPkt *>(msg);
			if((int)membersCH.size()<MAX_CLUSTER_SIZE)
				membersCH.push_back(m->getSrcAddr());

			//If I don't receive any other 'join' message after 'waitingTime' seconds -> TDMA schedule creation
			if(timerNextStage != NULL)
				cancelAndDelete(timerNextStage);

			//I must wait 'waitingTime' seconds to go to the next stage
			timerNextStage = new cMessage("Timer for creating the TDMA schedule", TIMER_CREATE_TDMA_SCHEDULE);
			scheduleAt(simTime() + waitingTime, timerNextStage);
		}
		delete msg;
	}
	break;

	case TDMA_SCHEDULE:
	{
		if(currentPhase != SCHEDULE_CREATION)
		{	//If we receive this message in a different phase -> Ignore it
			EV << "ERROR: 'TDMA schedule' message received during phase " << currentPhase << endl;
		}

		//TODO: Data Transmission Phase Eliminated
//		else{
//			//I must check if the message comes from my CH.
//			//If yes, I must wait till my slot to send data
//			LEACH_TDMA_Schedule* sched_pkt = static_cast<LEACH_TDMA_Schedule*>(msg);
//
//			if(sched_pkt->getSrcAddr()==myCH)
//			{
//				int my_slot=-1;
//
//				for(int i=0; i<(int)sched_pkt->getScheduleArraySize() && my_slot==-1; i++) {
//					if(sched_pkt->getSchedule(i)==myNetwAddr)
//						my_slot = i;
//				}
//
//				//If my slot, just send data NOW
//				if(my_slot==0 && packetToSend != NULL) {
//					dataTransmission();
//					this->switchHostState(HostState::SLEEP);
//				}
//				//If not, wait my slot and unplug the receiver
//				else
//					if(my_slot > 0){
//						this->switchHostState(HostState::SLEEP);
//						timerNextStage = new cMessage("Timer for waiting to my slot to send data", TIMER_SEND_DATA);
//						scheduleAt(simTime() + my_slot*slotTime, timerNextStage);
//					}
//
//				//If I don't have slot (my_slot == -1), do nothing = wait for the next round
//			}
//		}
		delete msg;
	}
	break;

	default:
	{
		if(currentPhase != DATA_TRANSMISSION)
		{	//If we receive this message in a different phase -> Ignore it
			EV << "ERROR: 'DATA' message received during phase " << currentPhase << endl;
			delete msg;
		}

		else {
			//If I don't know the message kind, it's supposed to be a DATA_MESSAGE.
			//If not, application layer will delete the message.
			//If I'm CH, I must store the message. If not, ignore it.

			if(lastRoundCH == currentRound)
				netwQueue.push_back(static_cast<NetwPkt *>(msg));

			else
				delete msg;
		}
	}
	break;

	}//end switch
}

void LEACH::handleSelfMsg(cMessage* msg){

	switch( msg->getKind() ){

	case TIMER_NEW_ROUND:
		//A new round starts: set timer for the next round and start with 'Advertisement Phase'
		currentRound++;
		delete msg;
		timerNextRound = new cMessage("New Round", TIMER_NEW_ROUND);
		scheduleAt(simTime() + roundTime, timerNextRound);

		this->switchHostState(HostState::ACTIVE);
		advertisementPhase();

		break;

	case TIMER_JOIN_CH:
		//Non-CHs shouldn't receive any more CH announce -> join to a CH
		//If I don't receive any CH announce, I'll wait for the next round
		if(distanceToCHs.size()>0)
			clusterSetUpPhase();

		delete msg;
		timerNextStage = NULL;
		break;

	case TIMER_CREATE_TDMA_SCHEDULE:
		//CH don't receive any more 'join' message -> TDMA schedule creation
		scheduleCreation();

		delete msg;
		timerNextStage= NULL;
		break;

	case TIMER_SEND_DATA:
		//Starts the data transmission for Non-CHs
		if(packetToSend != NULL){
			this->switchHostState(HostState::ACTIVE);
			dataTransmission();
			this->switchHostState(HostState::SLEEP);
		}

		delete msg;
		timerNextStage = NULL;
		break;
	}
}


//This packet contains data from the upper layer,
//so we must wait to our sending slot to send the data to the BS.

void LEACH::handleUpperMsg(cMessage* msg)
{
	if(packetToSend != NULL)
		delete packetToSend;
	packetToSend = encapsMsg(static_cast<cPacket*>(msg));
	EV 	<< "packet stored" << endl;
}

/**
 * Encapsulates the received ApplPkt into a NetwPkt and set all needed
 * header fields.
 **/

NetwPkt* LEACH::encapsMsg(cPacket *appPkt) {
//    int macAddr;
//    int netwAddr;

    EV <<"in encaps...\n";

    NetwPkt *pkt = new NetwPkt(appPkt->getName(), appPkt->getKind());
    pkt->setBitLength(headerLength);

    NetwControlInfo* cInfo = dynamic_cast<NetwControlInfo*>(appPkt->removeControlInfo());

//    if(cInfo == 0){
//	EV << "warning: Application layer did not specifiy a destination L3 address\n"
//	   << "\tusing broadcast address instead\n";
//	netwAddr = L3BROADCAST;
//    } else {
//	EV <<"CInfo removed, netw addr="<< cInfo->getNetwAddr()<<endl;
//        netwAddr = cInfo->getNetwAddr();
//	delete cInfo;
//    }

    pkt->setSrcAddr(myNetwAddr);
//    pkt->setDestAddr(myCH);
//    EV << " netw "<< myNetwAddr << " sending packet" <<endl;

//    switch(netwAddr) {
//
//    case L3BROADCAST:
//    	EV << "sendDown: nHop=L3BROADCAST -> message has to be broadcasted"
//    	   << " -> set destMac=L2BROADCAST\n";
//    	macAddr = L2BROADCAST;
//    	break;
//    case L3BS:
//    	EV << "sendDown: nHop=L3BS -> message for the BS"
//    	   << " -> set destMac=L2BS\n";
//    	macAddr = L2BS;
//    	break;
//    default:
//    	EV <<"sendDown: get the MAC address\n";
//    	macAddr = arp->getMacAddr(netwAddr);
//    	break;
//    }
//    pkt->setControlInfo(new MacControlInfo(macAddr));

    //encapsulate the application packet
    pkt->encapsulate(appPkt);
    EV <<" pkt encapsulated\n";
    return pkt;
}

void LEACH::handleLowerControl(cMessage* msg)
{
	switch (msg->getKind())
	{
	case BaseMacLayer::PACKET_DROPPED:
		droppedMsgs++;
		delete msg;
		break;
	default:
		EV << "LEACH does not handle control messages called " << msg->getName() << endl;
		delete msg;
	}
}

void LEACH::finish()
{
	recordScalar("dropped MAC messages", droppedMsgs);
}

