#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <libciomr/libciomr.h>
#include <libint/libint.h>
#include <libchkpt/chkpt.h>
#include <libiwl/iwl.h>
#include <psifiles.h>
#include <libqt/qt.h>

FILE *infile, *outfile;
char *psi_file_prefix;

void init_io(int argc, char *argv[]);
void exit_io(void);
void title(void);

int main(int argc, char *argv[])
{
  int iter, s, t, A, k, l, m, inew, iold, max;
  int i, j, ij, am, atom, shell_length, offset, stat;
  int nirreps, nao, nmo, nso, natom, nshell, noei, nocc, errcod, nfzc;
  int *stype, *snuc, *aostart, *aostop, *ao2atom, *l_length;
  int *clsdpi, *openpi, *orbspi, *dummy, *order, *frdocc;
  double *ss, **S, **scf, **u, **Ctmp, **C, *evals;

  int *orb_order, *orb_boolean, puream;
  double P, PiiA, Pst, Pss, Ptt, Ast, Bst, AB;
  double Uss, Utt, Ust, Uts, Cks, Ckt, **U, **V, **VV, **F;
  double cos4a, alpha, alphamax, alphalast, conv, norm;
  int print;

  alphalast = 1.0;

  init_io(argc, argv);
  title();

  chkpt_init(PSIO_OPEN_OLD);
  nao = chkpt_rd_nao();
  nmo = chkpt_rd_nmo();
  nso = chkpt_rd_nso();
  natom = chkpt_rd_natom();
  nshell = chkpt_rd_nshell();
  stype = chkpt_rd_stype();
  snuc = chkpt_rd_snuc();
  u = chkpt_rd_usotao();
  nirreps = chkpt_rd_nirreps();
  clsdpi = chkpt_rd_clsdpi();
  openpi = chkpt_rd_openpi();
  orbspi = chkpt_rd_orbspi();
  C = chkpt_rd_scf();
  evals = chkpt_rd_evals();
  chkpt_close();

  /* A couple of error traps */
  if(nirreps != 1) {
    fprintf(outfile, "\n\tError: localization is only valid in C1 symmetry!\n");
    exit(PSI_RETURN_FAILURE);
  }
  if(openpi[0]) {
    fprintf(outfile, "\n\tError: localization available for closed-shells only!\n");
    exit(PSI_RETURN_FAILURE);
  }

  /* Frozen orbital info */
  frdocc = get_frzcpi();
  nfzc = frdocc[0];
  free(frdocc);

  /* Compute the length of each AM block */
  ip_boolean("PUREAM", &(puream), 0);
  l_length = init_int_array(LIBINT_MAX_AM);
  l_length[0] = 1;
  for(l=1; l < (LIBINT_MAX_AM); l++) {
    if(puream) l_length[l] = 2 * l + 1;
    else l_length[l] = l_length[l-1] + l + 1;
  }

  /* Set up the atom->AO and AO->atom lookup arrays */
  aostart = init_int_array(natom);
  aostop = init_int_array(natom);
  for(i=0,atom=-1,offset=0; i < nshell; i++) {
    am = stype[i] - 1;
    shell_length = l_length[am];

    if(atom != snuc[i]-1) {
      if(atom != -1) aostop[atom] = offset-1;
      atom = snuc[i]-1;
      aostart[atom] = offset;
    }

    offset += shell_length;
  }
  aostop[atom] = offset-1;

  ao2atom = init_int_array(nso);
  for(i=0; i < natom; i++)
    for(j=aostart[i]; j <= aostop[i]; j++) ao2atom[j] = i;

  /* Get the overlap integrals -- these should be identical to AO S */
  noei = nso*(nso+1)/2;
  ss = init_array(noei);
  stat = iwl_rdone(PSIF_OEI,PSIF_SO_S,ss,noei,0,0,outfile);
  S = block_matrix(nso,nso);
  for(i=0,ij=0; i < nso; i++)
    for(j=0; j <= i; j++,ij++) {
      S[i][j] = S[j][i] = ss[ij];
    }
  free(ss);

  /* Compute nocc --- closed-shells only */
  for(i=0,nocc=0; i < nirreps; i++) nocc += clsdpi[i];

  fprintf(outfile, "\tNumber of doubly occupied orbitals: %d\n\n", nocc);

  fprintf(outfile, "\tIter     Pop. Localization   Max. Rotation Angle       Conv\n");
  fprintf(outfile, "\t------------------------------------------------------------\n");

  U = block_matrix(nocc, nocc);
  V = block_matrix(nocc, nocc);
  for(i=0; i < nocc; i++) V[i][i] = 1.0;
  VV = block_matrix(nocc, nocc);

  for(iter=0; iter < 100; iter++) {

    P = 0.0;
    for(i=nfzc; i < nocc; i++) {
      for(A=0; A < natom; A++) {
	PiiA = 0.0;

	for(l=aostart[A]; l <= aostop[A]; l++)
	  for(k=0; k < nso; k++) 
	    PiiA += C[k][i] * C[l][i] * S[k][l];

	P += PiiA * PiiA;
      }
    }

    /* Compute 2x2 rotations for Pipek-Mezey localization */
    alphamax = 0.0;
    for(s=nfzc; s < nocc; s++) {
      for(t=nfzc; t < s; t++) {

	Ast = Bst = 0.0;

	for(A=0; A < natom; A++) {

	  Pst = Pss = Ptt = 0.0;
	      
	  for(l=aostart[A]; l <= aostop[A]; l++) {
	    for(k=0; k < nso; k++) {
	      Pst += 0.5 * (C[k][s] * C[l][t] +
			    C[l][s] * C[k][t]) * S[k][l];

	      Pss += C[k][s] * C[l][s] * S[k][l];

	      Ptt += C[k][t] * C[l][t] * S[k][l];
	    }
	  }

	  Ast += Pst * Pst - 0.25 * (Pss - Ptt) * (Pss - Ptt);
	  Bst += Pst * (Pss - Ptt);

	} /* A-loop */

	/* Compute the rotation angle */
	AB = Ast * Ast + Bst * Bst;
	if(fabs(AB) > 0.0) {
	  cos4a = -Ast/sqrt(AB);
	  alpha = 0.25 * acos(cos4a) * (Bst > 0 ? 1 : -1);
	}
	else alpha = 0.0;

	/* Keep up with the maximum 2x2 rotation angle */
	alphamax = (fabs(alpha) > alphamax ? alpha : alphamax);

	Uss = cos(alpha);
	Utt = cos(alpha);
	Ust = sin(alpha);
	Uts = -Ust;

	/* Now do the rotation */
	for(k=0; k < nso; k++) {
	  Cks = C[k][s];
	  Ckt = C[k][t];
	  C[k][s] = Uss * Cks + Ust * Ckt;
	  C[k][t] = Uts * Cks + Utt * Ckt;
	}

	zero_mat(U, nocc, nocc);
	for(i=0; i < nocc; i++) U[i][i] = 1.0;

	U[s][s] = Uss;
	U[t][t] = Utt;
	U[s][t] = Ust;
	U[t][s] = Uts;

	zero_mat(VV, nocc, nocc);
	for(i=0; i < nocc; i++) {
	  for(j=0; j < nocc; j++) {
	    for(k=0; k < nocc; k++) {
	      VV[i][j] += V[i][k] * U[j][k];
	    }
	  }
	}

	for(i=0; i < nocc; i++)
	  for(j=0; j < nocc; j++)
	    V[i][j] = VV[i][j];
	      
      } /* t-loop */
    } /* s-loop */

    conv = fabs(alphamax) - fabs(alphalast);
    fprintf(outfile, "\t%4d  %20.10f  %20.10f  %4.3e\n", iter, P, alphamax, conv);
    if((iter > 2) && ((fabs(conv) < 1e-12) || alphamax == 0.0)) break;
    alphalast = alphamax;

    fflush(outfile);
      
  } /* iter-loop */

  /*  print_mat(V, nocc, nocc, outfile);  */

  /* Transform occupied orbital eigenvalues */
  F = block_matrix(nocc, nocc);
  for(i=0; i < nocc; i++)
    for(j=0; j < nocc; j++)
      for(k=0; k < nocc; k++) 
	F[i][j] += V[k][i] * evals[k] * V[k][j];

  /*
  fprintf(outfile, "\nTransformed Orbital Energies:\n");
  print_mat(F, nocc, nocc, outfile);
  */

  /* Compute a reordering array based on the diagonal elements of F */
  orb_order = init_int_array(nocc);
  orb_boolean = init_int_array(nocc);
  for(i=0; i < nocc; i++) { orb_order[i] = 0;  orb_boolean[i] = 0; }

  for(i=0,max=0; i < nocc; i++) /* First, find the overall maximum */
    if(fabs(F[i][i]) > fabs(F[max][max])) max = i;

  orb_order[0] = max;  orb_boolean[max] = 1;

  for(i=1; i < nocc; i++) {
    max = 0;
    while(orb_boolean[max]) max++; /* Find an unused max */
    for(j=0; j < nocc; j++) 
      if((fabs(F[j][j]) >= fabs(F[max][max])) && !orb_boolean[j]) max = j;
    orb_order[i] = max; orb_boolean[max] = 1;
  }

  /*
  for(i=0; i < nocc; i++) fprintf(outfile, "%d %d\n", i, orb_order[i]);
  */

 /*
  fprintf(outfile, "\n\tPipek-Mezey Localized MO's (before sort):\n");
  print_mat(C, nso, nmo, outfile);
  */

  /* Now reorder the localized MO's according to F */
  Ctmp = block_matrix(nso,nocc);
  for(i=0; i < nocc; i++)
    for(j=0; j < nso; j++) Ctmp[j][i] = C[j][i];

  for(i=0; i < nocc; i++) {
    iold = orb_order[i];
    for(j=0; j < nso; j++) C[j][i] = Ctmp[j][iold];
    evals[i] = F[iold][iold];
  }
  free_block(Ctmp);

  print = 0;
  errcod = ip_boolean("PRINT_MOS", &(print), 0);
  if(print) {
    fprintf(outfile, "\n\tPipek-Mezey Localized MO's (after sort):\n");
    print_mat(C, nao, nmo, outfile);
  }

  /* Check MO normalization */
  /*
  for(i=0; i < nmo; i++) {
    norm = 0.0;
    for(j=0; j < nao; j++) 
      for(k=0; k < nao; k++) {
	norm += C[j][i] * C[k][i] * S[j][k];
      }

    fprintf(outfile, "norm[%d] = %20.10f\n", i, norm);
  }
  */

  /* Write the new MO's to chkpt */
  chkpt_init(PSIO_OPEN_OLD);
  chkpt_wt_scf(C);
  chkpt_close();

  free_block(C);
  free(evals);

  fprintf(outfile, "\n\tLocalization of occupied orbitals complete.\n");

  exit_io();
  exit(PSI_RETURN_SUCCESS);
}

void init_io(int argc, char * argv[])
{
  extern char *gprgid();
  char *progid;

  progid = (char *) malloc(strlen(gprgid())+2);
  sprintf(progid, ":%s",gprgid());

  psi_start(argc-1,argv+1,0);
  ip_cwk_add(":INPUT"); /* for checking puream keyword */
  ip_cwk_add(progid);
  free(progid);
  tstart(outfile);

  psio_init();
}

void title(void)
{
  fprintf(outfile, "\n");
  fprintf(outfile, "\t\t\t**************************\n");
  fprintf(outfile, "\t\t\t*                        *\n");
  fprintf(outfile, "\t\t\t*          LOCAL         *\n");
  fprintf(outfile, "\t\t\t*                        *\n");
  fprintf(outfile, "\t\t\t**************************\n");
  fprintf(outfile, "\n");
}

void exit_io(void)
{
  psio_done();
  tstop(outfile);
  psi_stop();
}

char *gprgid()
{
   char *prgid = "LOCAL";

   return(prgid);
}
