/*
************************************************************************
*
*   opt.C -
*
*   Copyright (c) 1995
*
*   ETH Zuerich
*   Institut fuer Molekularbiologie und Biophysik
*   ETH-Hoenggerberg
*   CH-8093 Zuerich
*
*   All Rights Reserved
*
*   Date of last modification : 95/09/15
*   Pathname of SCCS file     : /export/home3/cb/garant-1.0/src/SCCS/s.opt.C
*   SCCS identification       : 1.3
*
************************************************************************
*/
/**************************************************************************/
/*                              opt.cc                                    */
/*                                                                        */
/* methods to optimize assignments described by the nmr diagram           */
/**************************************************************************/

#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "log.h"
#include "opt.h"
#include "choice.h"
#include "peak.h"
#include "erPeak.h"
#include "specTypes.h"
#include "envir.h"

/* instances of choice to save some memory */
ChoicePeak  choice;
ChoicePeak  choiceDegen;

OptPar::OptPar(float value, float variance, float l, float u)
{
   low = l;
   up = u;
   v = value;
   var = variance;
}

Optimizer::Optimizer(NMRdia &dia): nmrdia(dia),
				   tempCo(0.2,0.0,0.2,0.5),
				   ratioUnassMult(0.2,0.0,1E-3,1E+2),
				   mutIt(0.5,40.0,0.0,0.95),
		     	           nrInitPoints(10.0,0.0,2.0,500.0)
{
   parPeak=0; parNr=0;
}

OptGen::OptGen(NMRdia &dia): Optimizer(dia) 
{
}

/* Stochastic selection from parents */
Peak *Optimizer::selectParents(ErPeak *erPeak, float *low, float *up,
			    Wert *ass, Wert *exp, int *flags, int nrAss)
{
   int i,j;
   Peak *peak;
   float prob;
   int res;
   float random;

   parPeak=0; parNr=0;
   choice.reset();
   choiceDegen.reset();
   for(i=0;i<nmrdia.nrActive;i++) {
      peak = erPeak->getAss(nmrdia.activePos[i]);
      if(peak != 0) {
         res = 1;
         for(j=0;j<peak->p->spec->type->dim;j++)
            if(flags[j]==ASSIGNED && (peak->w[j] < low[j] || peak->w[j] > up[j]))
               res = 0;
         if(res) {
            prob = peak->calcPSelect(ass,exp,flags);
            if(peak->p->nrAss==0) {
               choice.insert(peak,prob,erPeak,nmrdia.activePos[i]);
            } else {
               choiceDegen.insert(peak,prob,erPeak,nmrdia.activePos[i]);
            }
         }
      }
   }
   random = drand48() * probSelDegen;
   peak = choice.select(parPeak,parNr);
   if(peak==0 && (nrAss == erPeak->spec->type->dim || random < erPeak->pSelect))
      peak= choiceDegen.select(parPeak,parNr);

   return peak;
}

/* Stochastic selection from equivalent peaks */
Peak *Optimizer::selectEquiv(LiEquiv &equiv, float *low, float *up,
			    Wert *ass, Wert *exp, int *flags,
//                            ErPeak *, int )
                            ErPeak *targetPeak, int nrAss)
{
   int i,j;
   Peak *peak;
   float prob;
   float random;
   int res;
   ErPeak *erPeak;

   parPeak=0; parNr=0;
   choice.reset();
   choiceDegen.reset();
   erPeak = equiv.start();
   while(erPeak != 0) {
      for(i=0;i<nmrdia.nrActive;i++) {
         peak = erPeak->getAss(nmrdia.activePos[i]);
         if(peak != 0) {
//            res =1;
            if(peak->p->fixed) res=0;
            else res=1;
            for(j=0;j<peak->p->spec->type->dim;j++)
               if(peak->w[j] < low[j] || peak->w[j] > up[j])
                  res = 0;
            if(res ) {
               prob = peak->calcPSelect(ass,exp,flags);
               if(peak->p->nrAss==0) {
                  choice.insert(peak,prob,erPeak,nmrdia.activePos[i]);
               } else {
                  choiceDegen.insert(peak,prob,erPeak,nmrdia.activePos[i]);
               }
            }
         }
      }
      erPeak = equiv.next();
   }

   random = drand48() * probSelDegen;
   peak = choice.select(parPeak,parNr);
   if(peak==0 &&
         (nrAss == targetPeak->spec->type->dim || random < targetPeak->pSelect))
      peak= choiceDegen.select(parPeak,parNr);

   return peak;
}

/* Stochastic selection from a given list of peaks */
Peak *Optimizer::selectList(Peak **list, int n, float *low, float *up,
			    Wert *ass, Wert *exp, int *flags, ErPeak *erPeak,
                            int nrAss)
{
   int i,j;
   Peak *peak;
   float prob;
   int res;
   float random;
   int notdegen=0;

   parPeak=0; parNr=0;
   choice.reset();
   choiceDegen.reset();
   for(i=0;i<n;i++) {
      peak = list[i];
      if(peak != 0) {
//         res = 1;
         if(peak->p->fixed) res=0;
         else res=1;
         for(j=0;j<peak->p->spec->type->dim;j++)
            if(peak->w[j] < low[j] || peak->w[j] > up[j]) res = 0;
         if(res) {
            prob = peak->calcPSelect(ass,exp,flags);
            if(peak->p->nrAss==0) {
               notdegen = choice.insert(peak,prob,0,0);
            } else if(!notdegen) {
               choiceDegen.insert(peak,prob,0,0);
            }
         }
      }
   }

   random = drand48() * probSelDegen;
   if(notdegen)
      return choice.select(parPeak,parNr);
   else if(nrAss == erPeak->spec->type->dim || random < erPeak->pSelect)
      return choiceDegen.select(parPeak,parNr);
   else
      return 0;
}

/* Stochastic selection from peaks within in given range */
Peak *Optimizer::selectRange(Spectrum *spec, float *low, float *up, 
                                 Wert *ass, Wert *exp, int nrAss, int *flags,
				ErPeak *erPeak)
{
   Peak *peak=0;
   float prob;
   int notdegen=0;
//   float random;

   parPeak=0; parNr=0;
   if(nrAss > erPeak->spec->type->dim-2) {
      choice.reset();
      choiceDegen.reset();
      peak = spec->peaks.start(low,up);
      while(peak != 0) {
         if(!peak->p->fixed) {
            prob = peak->calcPSelect(ass,exp,flags);
            if(peak->p->nrAss==0) {
               choice.insert(peak,prob,0,0);
               notdegen=1;
            } else if(!notdegen) {
               choiceDegen.insert(peak,prob,0,0);
            }
         }
         peak = spec->peaks.next();
      }
      if(notdegen)
         peak = choice.select(parPeak,parNr);
      else
         peak = choiceDegen.select(parPeak,parNr);
   } else {
      peak = spec->peaks.selectUnusedClose(exp,low,up);
   }
   return peak;
}

void Optimizer::initList()
{
   ErPeak *peak;
   Iterator i;

   liFixed.reset();
   liAss.reset();
   liReject.reset();
   liUnass.reset();

   peak = (ErPeak *)nmrdia.liErPeak.start(i);
   while(peak != 0) {
     peak->inRejectList = FALSE;
     peak->inInitList = FALSE;
     if(peak->type == SUBFRAG) {
     } else if(peak->status == P_FIXED) {
        liFixed.add(peak);
     } else if(peak->status == P_ASS) {
        liAss.addrandom(peak);
     } else if(peak->status == P_UNASS) {
        liUnass.addrandom(peak);
     } else if(peak->status == P_REJECTED) {
        peak->inRejectList = TRUE;
        liReject.addrandom(peak);
     }
     peak = (ErPeak *)nmrdia.liErPeak.next(i);
   }
}

void Optimizer::initAss()
{
   ErPeak *erPeak;
   Peak *peak;

   erPeak = getUnass();
   while(erPeak != 0) {
      peak = matchingInitPeak(erPeak);
      if(peak != 0 && accept(erPeak, peak) ) {
         acceptAss(erPeak, peak);
      } else {
         rejectAss(erPeak, peak);
      }
      erPeak = getUnass();
   }

   erPeak = getReject();
   while(erPeak != 0) {
      erPeak->status = P_UNASS;
      erPeak->inRejectList = FALSE;
      liUnass.add(erPeak);
      erPeak = getReject();
   }
}

void Optimizer::recursiveInit()
{
   Peak *peak;
   ErPeak *erPeak;
   ErPeak *neighbour;
   int i1,i2;
   int cnt;

   erPeak = (ErPeak *)liInit.get();
   while(erPeak != 0 && erPeak->status != P_UNASS)
      erPeak = (ErPeak *)liInit.get();
   while(erPeak != 0) {
      peak = matchingInitPeak(erPeak);
      if(peak != 0 && accept(erPeak, peak) ) {
         acceptAss(erPeak, peak);
         neighbour = erPeak->startNeighbour(i1,i2);
         cnt=0;
         while(neighbour != 0) {
            if(neighbour->status == P_UNASS && !neighbour->inInitList) {
               cnt++;
               neighbour->inInitList=TRUE;
	       liInit.addrandom(neighbour,cnt);
            }
            neighbour = erPeak->nextNeighbour(i1,i2);
         }
      } else {
         rejectAss(erPeak, peak);
      }
      erPeak = (ErPeak *)liInit.get();
      while(erPeak != 0 && erPeak->status != P_UNASS)
         erPeak = (ErPeak *)liInit.get();
   }
}

void Optimizer::opt(int maxIt)
{
   ErPeak *erPeak;
   Peak *peak;
   int flags[DIMENSION];
   int acceptMatch=0, rejectMatch=0, adaptPos=0, rejectPos=0;
   Log log;
   char str[MAXLINE];
   int i;
   int finish=0;

   for(i=0;i<DIMENSION+1;i++) {cntChanged[i]=0; }
   it = 0;
   probSelDegen = 10000.0;
   initList();
   initAss();
   probSelDegen = 3.0;
   while(!finish) {
      if(termCond() || it >= maxIt) finish=1;
      it++;
      erPeak = getUnass();
      while(erPeak != 0) {
         peak = matchingPeak(erPeak);
         if(peak != 0 && accept(erPeak, peak) ) {
            acceptAss(erPeak, peak);
            acceptMatch++;
         } else {
            rejectAss(erPeak, peak);
            rejectMatch++;
         }
         erPeak = getUnass();
      }

      if(!finish) {
         erPeak = getImprove();
         if(erPeak != 0) {
            peak = possPeak(erPeak);
            if(peak != 0 && adapt(erPeak,peak,flags)) {
               adaptCos(erPeak,peak,flags);
               adaptPos++;
            } else {
               keepAss(erPeak);
               rejectPos++;
           }
         }
      }
   }

   sprintf(str,
      "\n... %d of %d peaks matched, %d of %d peaks adapted:",
       acceptMatch, acceptMatch+rejectMatch,
       adaptPos, adaptPos+rejectPos);
   log.w(str);
   for(i=0;i<DIMENSION+1;i++) {
      sprintf(str,"\n    %4d times %d co's changed",cntChanged[i],i);
      log.w(str);
   }
}

ErPeak *Optimizer::getUnass()
{
   ErPeak *p;
   p = (ErPeak *)liUnass.get();
   while(p != 0 && p->status != P_UNASS) p = (ErPeak *)liUnass.get();
   return p;
}

ErPeak *Optimizer::getReject()
{
   ErPeak *p;
   static float factor = 1.0;
   int nrIt;
   float thresh;

   thresh = factor * drand48();
   p = (ErPeak *)liReject.get();
   while(p != 0 && p->status != P_REJECTED) {
      p->inRejectList = FALSE;
      p = (ErPeak *)liReject.get();
   }
   nrIt = 0;
   while(p != 0 && (p->pSelect < thresh || p->isPresent())  && nrIt < 100){
      liReject.add(p);
      p = (ErPeak *)liReject.get();
      while(p != 0 && p->status != P_REJECTED) {
         p->inRejectList = FALSE;
         p = (ErPeak *)liReject.get();
      }
      nrIt++; thresh *= 0.95;
   }
   if(nrIt > 10) factor *= 0.95;
   else if(nrIt < 3 && p != 0) factor *= 1.05;

   if(p!=0) {
      p->inRejectList = FALSE;
   }

   return p;
}

ErPeak *Optimizer::getUnassInit()
{
   ErPeak *p;

   p = (ErPeak *)liUnass.get();
   while(p != 0 && p->status != P_UNASS) p = (ErPeak *)liUnass.get();
   return p;
}

ErPeak *Optimizer::getAssDegen()
{
   ErPeak *p;
   static float factor = 1.0;
   int nrIt;
   float thresh;

   thresh = factor * drand48();

   p = (ErPeak *)liAss.get();
   while(p != 0 && p->status != P_ASS) p = (ErPeak *)liAss.get();
   nrIt = 0;
   while(p != 0 && nrIt < 100 &&
         ((p->getAss())->p->nrAss == (p->getAss())->p->nrIntraDegen(p)+1 ||
           p->pSelect < thresh
         )) {
      liAss.add(p);
      p = (ErPeak *)liAss.get();
      while(p != 0 && p->status != P_ASS) p = (ErPeak *)liAss.get();
      nrIt++; thresh *= 0.8;
   }
   if(nrIt > 10) factor *= 0.8;
   else if(nrIt < 3 && p != 0) factor *= 1.2;
   if(nrIt == 100) p = 0;

   return p;
}

ErPeak *Optimizer::getImprove()
{
   float randomnr;
   ErPeak *p=0;

   randomnr = drand48() * (ratioUnassMult.v + 1/ratioUnassMult.v);
   if(randomnr < ratioUnassMult.v) p = getAssDegen();
   if(!p) p = getReject();

   return p;
}

Peak *Optimizer::matchingPeak(ErPeak *erPeak)
{
   float low[DIMENSION];	/* lower bound of allowed range            */
   float up[DIMENSION];		/* upper bound of allowed range            */
   Wert ass[DIMENSION];		/* assigned frequencies                    */
   Wert exp[DIMENSION];		/* expected frequencies                    */
   int flag[DIMENSION];		/* are the coherences assigned ?           */
   int nrAss;			/* number of assigned frequencies          */
   Peak *peak = 0;

   erPeak->statCo(flag);
   erPeak->expFreq(exp);
   erPeak->assFreq(ass);
   nrAss = erPeak->range(flag,low,up);
   if(!erPeak->testDiagUnDef(flag)) {	/* no undefined diagonal  */
					/* coherences             */
      peak = selectRange(erPeak->spec, low, up, ass, exp,nrAss,flag,erPeak);
   }

   return peak;
}

int Optimizer::testCos(ErPeak *erPeak,int *status)
{
   float scoreCo[DIMENSION];	/* score of the coherences                 */
   float scoreCoDev[DIMENSION];	/* mean deviation from score               */
   float scoreCoMean[DIMENSION];/* average score for this coherence        */
   int i;
   float value;
   int nrSwitched=0;
   static float avg=0.0;
   float offset;

   erPeak->updateScoreCo(scoreCo,scoreCoMean,scoreCoDev);
   for(i=0;i<erPeak->spec->type->dim;i++) {
      scoreCo[i] -= scoreCoMean[i];
      scoreCo[i] /= scoreCoDev[i];
         if(avg<scoreCo[i]) {
            offset=0.05*tempCo.v;
         } else {
            offset=0.05 - 0.05*tempCo.v;
         }
         avg = (1.0-offset) * avg + offset*scoreCo[i];
   }
   for(i=0;i<erPeak->spec->type->dim;i++) {
      value = scoreCo[i] - avg;
      if(status[i] != UNASSIGNED && value<0  &&
          !erPeak->getCo(i)->fixed) {
         status[i] = UNASSIGNED;
         nrSwitched++;
      }
   }
   cntChanged[nrSwitched]++;
   return nrSwitched;
}

Peak *Optimizer::possPeak(ErPeak *erPeak)
{
   float low[DIMENSION];	/* lower bound of allowed range            */
   float up[DIMENSION];		/* upper bound of allowed range            */
   Wert ass[DIMENSION];		/* assigned frequencies                    */
   Wert exp[DIMENSION];		/* expected frequencies                    */
   int flag[DIMENSION];		/* are the coherences assigned ?           */
   int searchflag[DIMENSION];	/* search in extende region = UNASSIGNED   */
   int nrAss;			/* number of assigned frequencies          */
   Peak *peak=0;

   erPeak->statCo(flag);
   erPeak->statCo(searchflag);
   if(testCos(erPeak,searchflag) > 0 && !erPeak->testDiagUnDef(searchflag)) {
      erPeak->expFreq(exp);
      erPeak->assFreq(ass);
      nrAss = erPeak->range(searchflag,low,up);
      peak = selectRange(erPeak->spec, low, up, ass, exp,nrAss,flag,erPeak);
   }

   return peak;
}

int Optimizer::accept(ErPeak *, Peak *)
{
   return 1;
}

int Optimizer::adapt(ErPeak *erPeak, Peak *peak, int *decision)
{
   int flag[DIMENSION];		/* are the coherences assigned ?           */
   Wert ass[DIMENSION];		/* assigned frequencies                    */
   int i;

   erPeak->statCo(flag);
   erPeak->assFreq(ass);
   for(i=0;i<erPeak->spec->type->dim;i++) {
      if(flag[i] == ASSIGNED &&
         (peak->w[i] < ass[i].x-ass[i].v || peak->w[i] > ass[i].x+ass[i].v))
      {
         decision[i] = RESET;
      } else {
         decision[i] = ADAPT;
      }
   }
   return 1;
}

void Optimizer::makeAss(ErPeak *erPeak, Peak *peak, int *decision)
{
   ErPeak *neighbour;
   Coherence *co;
   int ip,i;

   for(i=0;i<erPeak->spec->type->dim;i++) {
      if(decision[i] == ADAPT && (!erPeak->getCo(i)->meanF.valid() ||
          !erPeak->getCo(i)->measF[erPeak->spec->nr].valid())) {
         co = erPeak->getCo(i);
         neighbour = co->startPeak(ip);
         while(neighbour != 0) {
            if(neighbour->status == P_REJECTED && neighbour != erPeak) {
               liUnass.add(neighbour);
               neighbour->status = P_UNASS;
            }
            neighbour = co->nextPeak(ip);
         }
      }
   }

   erPeak->makeAss(peak,decision,parPeak,parNr);
   liAss.add(erPeak);
}


void Optimizer::acceptAss(ErPeak *erPeak, Peak *peak)
{
   int i;
   int decision[DIMENSION];

   for(i=0;i<erPeak->spec->type->dim;i++) decision[i] = ADAPT;
   makeAss(erPeak,peak,decision);
}

void Optimizer::resetPeaks(ErPeak *erPeak, int *decision)
{
   int i1,i2;
   ErPeak *neighbour;

   erPeak->resetAss();
   neighbour = erPeak->startNeighbour(i1,i2,decision);
   while(neighbour != 0) {
      if(neighbour->status != P_UNASS && neighbour->status != P_FIXED) {
         liUnass.add(neighbour);
         neighbour->resetAss();
      }
      neighbour = erPeak->nextNeighbour(i1,i2,decision);
   }
}

void Optimizer::adaptCos(ErPeak *erPeak, Peak *peak, int *decision)
{
   resetPeaks(erPeak,decision);
   makeAss(erPeak,peak,decision);
}

void Optimizer::rejectAss(ErPeak *erPeak, Peak *)
{
   erPeak->rejectAss();
   liReject.add(erPeak);
   erPeak->inRejectList = TRUE;
}

void Optimizer::keepAss(ErPeak *erPeak)
{
   if(erPeak->status==P_REJECTED) {
      liUnass.add(erPeak);
      erPeak->status=P_UNASS;
      erPeak->inRejectList = FALSE;
   } else if(erPeak->status == P_ASS) {
      liAss.add(erPeak);
   } else {
      printf("\ndebugging Optimizer::keepAss\n");
   }
}

void Optimizer::printInfo()
{
   Log log;

   log.w("\n... local optimizer: ");
}

void OptGen::evolOpt(int optIt, Envir &, float tempS)
{
   tempSelect=tempS;

   /* optimize */
   nmrdia.resetAss();
   printInfo();
   opt(optIt);
   /* score assignment */
   nmrdia.scoreGlobal();
}

void OptGen::initAss()
{
   ErPeak *erPeak;
   ErPeak *neighbour;
   int i1,i2;
   int cnt;
   int i;

   /* initialization starting from peaks with fixed assignment */
   liInit.reset();
   erPeak = (ErPeak *)liFixed.start();
   cnt=0;
   while(erPeak != 0) {
      neighbour = erPeak->startNeighbour(i1,i2);
      while(neighbour != 0) {
         if(neighbour->status == P_UNASS && !neighbour->inInitList) {
             cnt++;
             neighbour->inInitList=TRUE;
	     liInit.addrandom(neighbour,cnt);
         }
         neighbour = erPeak->nextNeighbour(i1,i2);
      }
      erPeak = (ErPeak *)liFixed.next();
   }
   recursiveInit();

   /* initialization starting from random peaks */
   erPeak = getUnass();
   while(erPeak != 0) {
      liInit.reset();
      for(i=0;i<nrInitPoints.v && erPeak != 0;i++) {
         erPeak->inInitList=TRUE;
         liInit.add(erPeak);
         erPeak = getUnass();
      }
      recursiveInit();
   }

   erPeak = getReject();
   while(erPeak != 0) {
      erPeak->status = P_UNASS;
      erPeak->inRejectList = FALSE;
      liUnass.add(erPeak);
      erPeak = getReject();
   }
}

Peak *OptGen::matchingPeak(ErPeak *erPeak)
{
   float low[DIMENSION];	/* lower bound of allowed range            */
   float up[DIMENSION];		/* upper bound of allowed range            */
   Wert ass[DIMENSION];		/* assigned frequencies                    */
   Wert expFreq[DIMENSION];	/* expected frequencies                    */
   int flag[DIMENSION];		/* are the coherences assigned ?           */
   int nrAss;			/* number of assigned frequencies          */
   Peak *peak=0;
   float random;

   erPeak->statCo(flag);
   erPeak->expFreq(expFreq);
   erPeak->assFreq(ass);
   nrAss = erPeak->range(flag,low,up);

   if(nrAss == erPeak->spec->type->dim) {
      peak = selectRange(erPeak->spec, low, up, 
                         ass,expFreq,nrAss,flag,erPeak);
   } else {
      random = drand48();
      random = exp(log(random) * tempSelect);
      peak = selectParents(erPeak,low,up,ass,expFreq,flag,nrAss);
      if(peak==0 && erPeak->liEquiv != 0 &&
         random<=0.5)
         peak = selectEquiv(*erPeak->liEquiv,
                            low,up,ass,expFreq,flag,erPeak,nrAss);
      if(peak==0 && !erPeak->testDiagUnDef(flag) && 
         random<=0.3)
         peak = selectRange(erPeak->spec, low,up,ass,
                            expFreq,nrAss,flag,erPeak);
   }
   return peak;
}

Peak *OptGen::matchingInitPeak(ErPeak *erPeak)
{
   return matchingPeak(erPeak);
}

Peak *OptGen::possPeak(ErPeak *erPeak)
{
   float low[DIMENSION];	/* lower bound of allowed range            */
   float up[DIMENSION];		/* upper bound of allowed range            */
   Wert ass[DIMENSION];		/* assigned frequencies                    */
   Wert expFreq[DIMENSION];	/* expected frequencies                    */
   int flag[DIMENSION];		/* are the coherences assigned ?           */
   int searchflag[DIMENSION];	/* search in extende region = UNASSIGNED   */
   int nrAss;			/* number of assigned frequencies          */
   Peak *peak=0;
   float random;

   erPeak->statCo(flag);
   erPeak->statCo(searchflag);
   if(testCos(erPeak,searchflag) > 0 && !erPeak->testDiagUnDef(searchflag)) {
      erPeak->expFreq(expFreq);
      erPeak->assFreq(ass);
      nrAss = erPeak->range(searchflag,low,up);
      random = drand48();
      random = exp(log(random) * tempSelect);
      peak = selectParents(erPeak,low,up,ass,expFreq,flag,nrAss);
      if(peak==0 && erPeak->liEquiv != 0 &&
         random<=0.5)
         peak = selectEquiv(*erPeak->liEquiv,
                            low,up,ass,expFreq,flag,erPeak,nrAss);
      if(peak==0 && !erPeak->testDiagUnDef(flag) && 
         random<=0.3)
         peak = selectRange(erPeak->spec, low,up,ass,
                            expFreq,nrAss,flag,erPeak);
   }

   return peak;
}

void OptGen::printInfo()
{
   Log log;

   log.w("\n... genetic optimizer: ");
}