#include "itensor/all.h"

using namespace itensor;

IQIndex
directSum(IQIndex const& i, IQIndex const& j, Args const& args = Args::global())
  {
  auto name = args.getString("IndexName","sum");
#ifdef DEBUG
  if( i.dir() != j.dir() ) Error("In directSum(IQIndex, IQIndex), input indices must have same arrow direction");
#endif
  auto siq = stdx::reserve_vector<IndexQN>(i.nindex()+j.nindex());
  for(auto iq1 : i)
      {
      auto s1 = Index(iq1.index.rawname(),iq1.m(),iq1.type());
      siq.emplace_back(s1,iq1.qn);
      }
  for(auto iq2 : j)
      {
      auto s2 = Index(iq2.index.rawname(),iq2.m(),iq2.type());
      siq.emplace_back(s2,iq2.qn);
      }
#ifdef DEBUG
  if(siq.empty()) Error("siq is empty in plussers");
#endif
  auto ij = IQIndex(name,std::move(siq),i.dir());
  return ij;
  }

IQIndex
directSumTensors(IQIndex const& i,
                 IQIndex const& j,
                 IQTensor& D1, IQTensor& D2,
                 Args const& args = Args::global())
  {
  auto ij = directSum(i,j,args);
  D1 = IQTensor(dag(i),ij);
  int n = 1;
  for(auto iq1 : i)
      {
      auto s1 = ij.index(n);
      auto D = Matrix(iq1.index.m(),s1.m());
      auto minsize = std::min(iq1.index.m(),s1.m());
      for(auto ii : range(minsize)) D(ii,ii) = 1.0;
      D1 += matrixTensor(std::move(D),iq1.index,s1);
      ++n;
      }
  D2 = IQTensor(dag(j),ij);
  for(auto iq2 : j)
      {
      auto s2 = ij.index(n);
      auto D = Matrix(iq2.index.m(),s2.m());
      auto minsize = std::min(iq2.index.m(),s2.m());
      for(auto jj : range(minsize)) D(jj,jj) = 1.0;
      D2 += matrixTensor(std::move(D),iq2.index,s2);
      ++n;
      }
  return ij;
  }

IQTensor
directSum(IQTensor const& A, IQTensor const& B,
          IQIndexSet const& I, IQIndexSet const& J,
          IQIndexSet& IJ, Args const& args = Args::global())
  {
  if( I.size() != J.size() ) Error("In directSum(IQTensor, IQTensor, ...), must sum equal number of indices");
  auto AD = A;
  auto BD = B;
  auto newinds = IQIndexSetBuilder(I.size());
  for( auto n : range1(I.size()) )
    {
    auto In = I.index(n);
    auto Jn = J.index(n);
    if( dir(A,In) != In.dir() ) In.dag();
    if( dir(B,Jn) != Jn.dir() ) Jn.dag();
    IQTensor D1, D2;
    auto IJn = directSumTensors(In,Jn,D1,D2,args);
    newinds.nextIndex(IJn);
    AD *= D1;
    BD *= D2;
    }
  IJ = newinds.build();
  auto C = AD+BD;
  return C;
  }

// Direct sum A and B over indices i on A and j on B
IQTensor
directSum(IQTensor const& A, IQTensor const& B,
          IQIndex const& i, IQIndex const& j,
          IQIndex& ij, Args const& args = Args::global())
  {
  IQIndexSet IJ;
  auto C = directSum(A,B,IQIndexSet(i),IQIndexSet(j),IJ,args);
  ij = IJ.index(1);
  return C;
  }

void
test()
  {
  auto a1 = Index("a1",2);
  auto a2 = Index("a2",2);
  auto a3 = Index("a3",2);
  auto a = IQIndex("a",
                   a1,QN(-1),
                   a2,QN(0),
                   a3,QN(+1));

  auto b1 = Index("b1",2);
  auto b2 = Index("b2",2);
  auto b3 = Index("b3",2);
  auto b = IQIndex("b",
                   b1,QN(-1),
                   b2,QN(0),
                   b3,QN(+1));

  auto i1 = Index("i1",2);
  auto i2 = Index("i2",2);
  auto i3 = Index("i2",2);
  auto i = IQIndex("i",
                   i1,QN(-1),
                   i2,QN(0),
                   i3,QN(+1));

  auto j1 = Index("j1",2);
  auto j2 = Index("j2",2);
  auto j3 = Index("j3",2);
  auto j = IQIndex("j",
                   j1,QN(-1),
                   j2,QN(0),
                   j3,QN(+1));

  auto A = randomTensor(QN(0),a,b,dag(i));
  auto B = randomTensor(QN(0),a,b,dag(j));

  PrintData(A);
  PrintData(B);

  // Version accepting an index on A and an index on B to be direct summed
  // Here we create a new IQTensor C with indices {a,b,i+j}
  // Indices that are not specified must be shared by A and B
  IQIndex ij;
  auto C = directSum(A,B,i,j,ij,{"IndexName=","i+j"});

  // Version accepting sets of indices on A and B to be direct summed
  // Here we create a new IQTensor C with indices {a,i+j,b+b}
  //IQIndexSet ij;
  //auto C = directSum(A,B,{i,b},{j,b},ij,{"IndexName=","sum"});

  PrintData(C);
  PrintData(ij);
  }

int
main()
  {
  test();
  return 0;
  }