/*
 * This file is a part of TiledArray.
 * Copyright (C) 2013  Virginia Tech
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU 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 General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <iomanip>
#include <tiledarray.h>
#include "input_data.h"

using namespace TiledArray;
using namespace TiledArray::expressions;

int main(int argc, char** argv) {
  // Initialize runtime
  TiledArray::World& world = TiledArray::initialize(argc, argv);

  std::string file_name = argv[1];

  // Open input file.
  std::ifstream input(file_name.c_str());

  if(! input.fail()) {

    // Read input data.

    if(world.rank() == 0)
      std::cout << "Reading input...";

    InputData data(input);
    input.close();

    if(world.rank() == 0)
      std::cout << " done.\nConstructing Fock tensors...";

    // Construct Fock tensor
    TiledArray::TSpArrayD f_a_oo = data.make_f(world, alpha, occ, occ);
    TiledArray::TSpArrayD f_a_vv = data.make_f(world, alpha, vir, vir);
    // Just make references to the data since the input is closed shell.
    TiledArray::TSpArrayD& f_b_oo = f_a_oo;
    TiledArray::TSpArrayD& f_b_vv = f_a_vv;

    // Fence to make sure Fock tensors are initialized on all nodes
    world.gop.fence();

    if(world.rank() == 0)
      std::cout << " done.\nConstructing v_ab tensors...";

    // Construct the integral tensors
    TiledArray::TSpArrayD v_ab_oooo = data.make_v_ab(world, occ, occ, occ, occ);
    TiledArray::TSpArrayD v_ab_vvoo = data.make_v_ab(world, vir, vir, occ, occ);
    TiledArray::TSpArrayD v_ab_oovv = data.make_v_ab(world, occ, occ, vir, vir);
    TiledArray::TSpArrayD v_ab_vovo = data.make_v_ab(world, vir, occ, vir, occ);
    TiledArray::TSpArrayD v_ab_ovov = data.make_v_ab(world, occ, vir, occ, vir);
    TiledArray::TSpArrayD v_ab_voov = data.make_v_ab(world, vir, occ, occ, vir);
    TiledArray::TSpArrayD v_ab_ovvo = data.make_v_ab(world, occ, vir, vir, occ);
    TiledArray::TSpArrayD v_ab_vvvv = data.make_v_ab(world, vir, vir, vir, vir);

    // Fence to make sure data on all nodes has been initialized
    world.gop.fence();

    if(world.rank() == 0)
      std::cout << " done.\nConstructing v_aa and v_bb tensors...";

    TiledArray::TSpArrayD v_aa_oooo;
    v_aa_oooo("i,j,k,l") = v_ab_oooo("i,j,k,l") - v_ab_oooo("i,j,l,k");
    TiledArray::TSpArrayD v_aa_vvoo;
    v_aa_vvoo("a,b,i,j") = v_ab_vvoo("a,b,i,j") - v_ab_vvoo("a,b,j,i");
    TiledArray::TSpArrayD v_aa_vovo;
    v_aa_vovo("a,i,b,j") = v_ab_vovo("a,i,b,j") - v_ab_voov("a,i,j,b");
    TiledArray::TSpArrayD v_aa_oovv;
    v_aa_oovv("i,j,a,b") = v_ab_oovv("i,j,a,b") - v_ab_oovv("i,j,b,a");
    TiledArray::TSpArrayD v_aa_vvvv;
    v_aa_vvvv("a,b,c,d") = v_ab_vvvv("a,b,c,d") - v_ab_vvvv("a,b,d,c");
    // Just make references to the data since the input is closed shell.
    TiledArray::TSpArrayD& v_bb_oooo = v_aa_oooo;
    TiledArray::TSpArrayD& v_bb_vvoo = v_aa_vvoo;
    TiledArray::TSpArrayD& v_bb_vovo = v_aa_vovo;
    TiledArray::TSpArrayD& v_bb_oovv = v_aa_oovv;
    TiledArray::TSpArrayD& v_bb_vvvv = v_aa_vvvv;

    // Fence again to make sure data all the integral tensors have been initialized
    world.gop.fence();

    if(world.rank() == 0)
      std::cout << " done.\n";

    TiledArray::TSpArrayD t_aa_vvoo(world, v_aa_vvoo.trange(), v_aa_vvoo.shape());
    for(auto it = t_aa_vvoo.range().begin(); it != t_aa_vvoo.range().end(); ++it)
      if(t_aa_vvoo.is_local(*it) && (! t_aa_vvoo.is_zero(*it)))
        t_aa_vvoo.set(*it, 0.0);

    TiledArray::TSpArrayD t_ab_vvoo(world, v_ab_vvoo.trange(), v_ab_vvoo.shape());
    for(auto it = t_ab_vvoo.range().begin(); it != t_ab_vvoo.range().end(); ++it)
      if(t_ab_vvoo.is_local(*it) && (! t_ab_vvoo.is_zero(*it)))
        t_ab_vvoo.set(*it, 0.0);

    TiledArray::TSpArrayD t_bb_vvoo(world, v_bb_vvoo.trange(), v_bb_vvoo.shape());
    for(auto it = t_bb_vvoo.range().begin(); it != t_bb_vvoo.range().end(); ++it)
      if(t_bb_vvoo.is_local(*it) && (! t_bb_vvoo.is_zero(*it)))
        t_bb_vvoo.set(*it, 0.0);

    TiledArray::TSpArrayD D_vvoo(world, v_ab_vvoo.trange(), v_ab_vvoo.shape());
    for(auto it = D_vvoo.range().begin(); it != D_vvoo.range().end(); ++it)
      if(D_vvoo.is_local(*it) && (! D_vvoo.is_zero(*it)))
        D_vvoo.set(*it, world.taskq.add(data, & InputData::make_D_vvoo_tile, D_vvoo.trange().make_tiles_range(*it)));


    world.gop.fence();

    data.clear();

    if(world.rank() == 0)
      std::cout << "Calculating t amplitudes...\n";


    double energy = 0.0;

    for(unsigned int i = 0ul; i < 100; ++i) {

      if(world.rank() == 0)
        std::cout << "Iteration " << i << "\n";

      TiledArray::TSpArrayD r_aa_vvoo;
      r_aa_vvoo("p1a,p2a,h1a,h2a") =
          v_aa_vvoo("p1a,p2a,h1a,h2a")
          -f_a_vv("p1a,p3a")*t_aa_vvoo("p2a,p3a,h1a,h2a")
          +f_a_vv("p2a,p3a")*t_aa_vvoo("p1a,p3a,h1a,h2a")
          +f_a_oo("h3a,h1a")*t_aa_vvoo("p1a,p2a,h2a,h3a")
          -f_a_oo("h3a,h2a")*t_aa_vvoo("p1a,p2a,h1a,h3a")
          +0.5*t_aa_vvoo("p3a,p4a,h1a,h2a")*v_aa_vvvv("p1a,p2a,p3a,p4a")
          +v_ab_voov("p1a,h3b,h1a,p3b")*t_ab_vvoo("p2a,p3b,h2a,h3b")
          -v_aa_vovo("p1a,h3a,p3a,h1a")*t_aa_vvoo("p2a,p3a,h2a,h3a")
          -v_ab_voov("p1a,h3b,h2a,p3b")*t_ab_vvoo("p2a,p3b,h1a,h3b")
          +v_aa_vovo("p1a,h3a,p3a,h2a")*t_aa_vvoo("p2a,p3a,h1a,h3a")
          -v_ab_voov("p2a,h3b,h1a,p3b")*t_ab_vvoo("p1a,p3b,h2a,h3b")
          +v_aa_vovo("p2a,h3a,p3a,h1a")*t_aa_vvoo("p1a,p3a,h2a,h3a")
          +v_ab_voov("p2a,h3b,h2a,p3b")*t_ab_vvoo("p1a,p3b,h1a,h3b")
          -v_aa_vovo("p2a,h3a,p3a,h2a")*t_aa_vvoo("p1a,p3a,h1a,h3a")
          +0.5*v_aa_oooo("h3a,h4a,h1a,h2a")*t_aa_vvoo("p1a,p2a,h3a,h4a")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p2a,p4b,h3a,h4b")*t_aa_vvoo("p1a,p3a,h1a,h2a")
          -0.5*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p2a,p4a,h3a,h4a")*t_aa_vvoo("p1a,p3a,h1a,h2a")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p1a,p4b,h3a,h4b")*t_aa_vvoo("p2a,p3a,h1a,h2a")
          +0.5*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p1a,p4a,h3a,h4a")*t_aa_vvoo("p2a,p3a,h1a,h2a")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h2a,h4b")*t_aa_vvoo("p1a,p2a,h1a,h3a")
          -0.5*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p3a,p4a,h2a,h4a")*t_aa_vvoo("p1a,p2a,h1a,h3a")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h1a,h4b")*t_aa_vvoo("p1a,p2a,h2a,h3a")
          -0.5*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p3a,p4a,h1a,h3a")*t_aa_vvoo("p1a,p2a,h2a,h4a")
          +0.25*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p3a,p4a,h1a,h2a")*t_aa_vvoo("p1a,p2a,h3a,h4a")
          +v_bb_oovv("h3b,h4b,p3b,p4b")*t_ab_vvoo("p1a,p3b,h1a,h3b")*t_ab_vvoo("p2a,p4b,h2a,h4b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p1a,p4b,h1a,h4b")*t_aa_vvoo("p2a,p3a,h2a,h3a")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_aa_vvoo("p1a,p3a,h1a,h3a")*t_ab_vvoo("p2a,p4b,h2a,h4b")
          +v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p1a,p3a,h1a,h3a")*t_aa_vvoo("p2a,p4a,h2a,h4a")
          -v_bb_oovv("h3b,h4b,p3b,p4b")*t_ab_vvoo("p2a,p3b,h1a,h3b")*t_ab_vvoo("p1a,p4b,h2a,h4b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p2a,p4b,h1a,h4b")*t_aa_vvoo("p1a,p3a,h2a,h3a")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_aa_vvoo("p2a,p3a,h1a,h3a")*t_ab_vvoo("p1a,p4b,h2a,h4b")
          -v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p2a,p3a,h1a,h3a")*t_aa_vvoo("p1a,p4a,h2a,h4a");

      world.gop.fence();

      TiledArray::TSpArrayD r_ab_vvoo;
      r_ab_vvoo("p1a,p2b,h1a,h2b") =
          v_ab_vvoo("p1a,p2b,h1a,h2b")
          +f_a_vv("p1a,p3a")*t_ab_vvoo("p3a,p2b,h1a,h2b")
          +f_b_vv("p2b,p3b")*t_ab_vvoo("p1a,p3b,h1a,h2b")
          -f_a_oo("h3a,h1a")*t_ab_vvoo("p1a,p2b,h3a,h2b")
          -f_b_oo("h3b,h2b")*t_ab_vvoo("p1a,p2b,h1a,h3b")
          +t_ab_vvoo("p3a,p4b,h1a,h2b")*v_ab_vvvv("p1a,p2b,p3a,p4b")
          +v_ab_voov("p1a,h3b,h1a,p3b")*t_bb_vvoo("p2b,p3b,h2b,h3b")
          -v_aa_vovo("p1a,h3a,p3a,h1a")*t_ab_vvoo("p3a,p2b,h3a,h2b")
          -v_ab_vovo("p1a,h3b,p3a,h2b")*t_ab_vvoo("p3a,p2b,h1a,h3b")
          -v_ab_ovov("h3a,p2b,h1a,p3b")*t_ab_vvoo("p1a,p3b,h3a,h2b")
          -v_bb_vovo("p2b,h3b,p3b,h2b")*t_ab_vvoo("p1a,p3b,h1a,h3b")
          +v_ab_ovvo("h3a,p2b,p3a,h2b")*t_aa_vvoo("p1a,p3a,h1a,h3a")
          +v_ab_oooo("h3a,h4b,h1a,h2b")*t_ab_vvoo("p1a,p2b,h3a,h4b")
          -0.5*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p2b,p4b,h3b,h4b")*t_ab_vvoo("p1a,p3b,h1a,h2b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p2b,h3a,h4b")*t_ab_vvoo("p1a,p4b,h1a,h2b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p1a,p4b,h3a,h4b")*t_ab_vvoo("p3a,p2b,h1a,h2b")
          -0.5*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p1a,p4a,h3a,h4a")*t_ab_vvoo("p3a,p2b,h1a,h2b")
          -0.5*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p3b,p4b,h2b,h4b")*t_ab_vvoo("p1a,p2b,h1a,h3b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h3a,h2b")*t_ab_vvoo("p1a,p2b,h1a,h4b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h1a,h4b")*t_ab_vvoo("p1a,p2b,h3a,h2b")
          +0.5*v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p3a,p4a,h1a,h3a")*t_ab_vvoo("p1a,p2b,h4a,h2b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h1a,h2b")*t_ab_vvoo("p1a,p2b,h3a,h4b")
          +v_bb_oovv("h3b,h4b,p3b,p4b")*t_ab_vvoo("p1a,p3b,h1a,h3b")*t_bb_vvoo("p2b,p4b,h2b,h4b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p1a,p4b,h1a,h4b")*t_ab_vvoo("p3a,p2b,h3a,h2b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_aa_vvoo("p1a,p3a,h1a,h3a")*t_bb_vvoo("p2b,p4b,h2b,h4b")
          +v_aa_oovv("h3a,h4a,p3a,p4a")*t_aa_vvoo("p1a,p3a,h1a,h3a")*t_ab_vvoo("p4a,p2b,h4a,h2b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p2b,h1a,h4b")*t_ab_vvoo("p1a,p4b,h3a,h2b");

      world.gop.fence();

      TiledArray::TSpArrayD r_bb_vvoo;
      r_bb_vvoo("p1b,p2b,h1b,h2b") =
          v_bb_vvoo("p1b,p2b,h1b,h2b")
          -f_b_vv("p1b,p3b")*t_bb_vvoo("p2b,p3b,h1b,h2b")
          +f_b_vv("p2b,p3b")*t_bb_vvoo("p1b,p3b,h1b,h2b")
          +f_b_oo("h3b,h1b")*t_bb_vvoo("p1b,p2b,h2b,h3b")
          -f_b_oo("h3b,h2b")*t_bb_vvoo("p1b,p2b,h1b,h3b")
          +0.5*t_bb_vvoo("p3b,p4b,h1b,h2b")*v_bb_vvvv("p1b,p2b,p3b,p4b")
          -v_bb_vovo("p1b,h3b,p3b,h1b")*t_bb_vvoo("p2b,p3b,h2b,h3b")
          +v_ab_ovvo("h3a,p1b,p3a,h1b")*t_ab_vvoo("p3a,p2b,h3a,h2b")
          +v_bb_vovo("p1b,h3b,p3b,h2b")*t_bb_vvoo("p2b,p3b,h1b,h3b")
          -v_ab_ovvo("h3a,p1b,p3a,h2b")*t_ab_vvoo("p3a,p2b,h3a,h1b")
          +v_bb_vovo("p2b,h3b,p3b,h1b")*t_bb_vvoo("p1b,p3b,h2b,h3b")
          -v_ab_ovvo("h3a,p2b,p3a,h1b")*t_ab_vvoo("p3a,p1b,h3a,h2b")
          -v_bb_vovo("p2b,h3b,p3b,h2b")*t_bb_vvoo("p1b,p3b,h1b,h3b")
          +v_ab_ovvo("h3a,p2b,p3a,h2b")*t_ab_vvoo("p3a,p1b,h3a,h1b")
          +0.5*v_bb_oooo("h3b,h4b,h1b,h2b")*t_bb_vvoo("p1b,p2b,h3b,h4b")
          -0.5*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p2b,p4b,h3b,h4b")*t_bb_vvoo("p1b,p3b,h1b,h2b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p2b,h3a,h4b")*t_bb_vvoo("p1b,p4b,h1b,h2b")
          +0.5*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p1b,p4b,h3b,h4b")*t_bb_vvoo("p2b,p3b,h1b,h2b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p1b,h3a,h4b")*t_bb_vvoo("p2b,p4b,h1b,h2b")
          -0.5*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p3b,p4b,h2b,h4b")*t_bb_vvoo("p1b,p2b,h1b,h3b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h3a,h2b")*t_bb_vvoo("p1b,p2b,h1b,h4b")
          -0.5*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p3b,p4b,h1b,h3b")*t_bb_vvoo("p1b,p2b,h2b,h4b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p4b,h3a,h1b")*t_bb_vvoo("p1b,p2b,h2b,h4b")
          +0.25*v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p3b,p4b,h1b,h2b")*t_bb_vvoo("p1b,p2b,h3b,h4b")
          +v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p1b,p3b,h1b,h3b")*t_bb_vvoo("p2b,p4b,h2b,h4b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p2b,h3a,h2b")*t_bb_vvoo("p1b,p4b,h1b,h4b")
          +v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p1b,h3a,h1b")*t_bb_vvoo("p2b,p4b,h2b,h4b")
          +v_aa_oovv("h3a,h4a,p3a,p4a")*t_ab_vvoo("p3a,p1b,h3a,h1b")*t_ab_vvoo("p4a,p2b,h4a,h2b")
          -v_bb_oovv("h3b,h4b,p3b,p4b")*t_bb_vvoo("p2b,p3b,h1b,h3b")*t_bb_vvoo("p1b,p4b,h2b,h4b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p1b,h3a,h2b")*t_bb_vvoo("p2b,p4b,h1b,h4b")
          -v_ab_oovv("h3a,h4b,p3a,p4b")*t_ab_vvoo("p3a,p2b,h3a,h1b")*t_bb_vvoo("p1b,p4b,h2b,h4b")
          -v_aa_oovv("h3a,h4a,p3a,p4a")*t_ab_vvoo("p3a,p2b,h3a,h1b")*t_ab_vvoo("p4a,p1b,h4a,h2b");


      world.gop.fence();

      t_aa_vvoo("a,b,i,j") =
          D_vvoo("a,b,i,j") * r_aa_vvoo("a,b,i,j")
          + t_aa_vvoo("a,b,i,j");

      t_ab_vvoo("a,b,i,j") =
          D_vvoo("a,b,i,j") * r_ab_vvoo("a,b,i,j")
          + t_ab_vvoo("a,b,i,j");

      t_bb_vvoo("a,b,i,j") =
          D_vvoo("a,b,i,j") * r_bb_vvoo("a,b,i,j")
          + t_bb_vvoo("a,b,i,j");

      const double error = (r_aa_vvoo("a,b,i,j")
          + r_ab_vvoo("a,b,i,j") + r_bb_vvoo("a,b,i,j")).norm();

      energy =
           0.25 * ( t_aa_vvoo("a,b,i,j").dot(v_aa_vvoo("a,b,i,j"))
                   + t_bb_vvoo("a,b,i,j").dot(v_bb_vvoo("a,b,i,j"))
                  )
          + t_ab_vvoo("a,b,i,j").dot(v_ab_vvoo("a,b,i,j"));

      world.gop.fence();

      if(world.rank() == 0)
        std::cout << " error  = "  << std::setprecision(12) << error << "\n"
                  << " energy = " << std::setprecision(12) << energy << "\n";

      if(error < 1.0e-10)
        break;
    }

    if(world.rank() == 0) {
      std::cout << "CCD energy = " << std::setprecision(12) << energy << "\n";
      std::cout << "Done!\n";
    }

  } else  {
    std::cout << "Unable to open file: " << file_name << "\n";
    // stop the madenss runtime
    TiledArray::finalize();
    return 1;
  }

  TiledArray::finalize();
	return 0;
}
