#include "./itensor/all.h"
using namespace itensor;
#include <iomanip>

void Sx_meas( MPS psi, SiteSet sites ) // Sx measurement and output
{
    int N = length(sites);

    for( int j = 1; j <= N; j++ )
    {
        psi.position(j);

        auto ket = psi(j);
        auto bra = dag(prime(ket,"Site"));
        auto Sxjop = sites.op("Sx",j);

        auto sxj = 2*eltC(bra*Sxjop*ket).real();

        std::cout << std::setprecision(3) << sxj << " ";
    }
}

MPO get_H( const SiteSet sites, double h, double hx )
{
    auto ampo = AutoMPO(sites);
    int N = length(sites);
    for (int i = 1; i < N; i++)
    {
        // produces correct GS
        ampo += -0.5, "S+", i, "S+", i+1;
        ampo += -0.5, "S-", i, "S-", i+1;
        ampo += -0.5, "S+", i, "S-", i+1;
        ampo += -0.5, "S-", i, "S+", i+1;

        // produces wrong GS
        //ampo += -1, "Sx", i, "Sx", i+1;
    }
    for (int i = 1; i <= N; i++)
    {
        ampo += -h, "Sz", i;
        ampo += -hx, "Sx", i;
    }

    MPO H_TI = toMPO(ampo);

    return H_TI;
}

int main(int argc, char **argv)
{
    int N = 50;
    double h = 0.5;

    auto sites = SpinHalf(N,{"ConserveQNs=",false});

    auto sweeps = Sweeps(10);
    sweeps.maxdim() = 1000;
    sweeps.cutoff() = 1E-10;

    std::vector<double> hxs = {1.0, 0.5, 0.4, 0.3, 0.2, 0.1, 0.01, 1e-3, 0};

    MPS psi_up = randomMPS(sites);
    double energy_up;

    for( double hxk : hxs )
    {
        auto H_TI_up = get_H( sites, h, hxk); // Construct H with smaller and smaller hx to end up in the symmetry-broken GS
        std::tie(energy_up,psi_up) = dmrg(H_TI_up,psi_up,sweeps,{"Quiet",true});
        psi_up.normalize();

        Sx_meas( psi_up, sites );
    }
}