import ranrot.*;

//
// System
// 

record System
(
    r: Pointer[RanrotBGen],
    n: Int,
    board: Vector[Int],
    rdiag: Vector[Int],
    ldiag: Vector[Int],
    E: Int,
);

overload System(n, r)
{
    var s = System();
    s.n = n;
    s.r = &r;

    resize(s.board, n);
    resize(s.rdiag, 2*n);
    resize(s.ldiag, 2*n);
    
    sample(s);
    
    return s;
}

sample(s: System)
{
    for(x in range(s.n))
        s.board[x] = x; // biased, valid initial board...
        
    // burn-in
    for (i in range(100))
    {
        var r1 = randomInt(s.r^, 0, s.n-1);
        var r2 = randomInt(s.r^, 0, s.n-1);
        swap(s.board[r1], s.board[r2]);
    }
        
    setDiags(s);
    s.E = computeEnergy(s);
}

var noRandom = true;

var cached = 0;
var cachedNum = 0;

fastRandomInt(r: RanrotBGen, min, max) Int
{
    if (true) return randomInt(r, min, max);
        
        cached += 101;
        if (cached>100000) cached = 0;
        
        return cached%max;
        
      /*  
    if (cached <= 0)
    {
        cachedNum = randomInt(r, min, 1024*1024*1024);
        cached = 400;
    }

    var r = cachedNum % max;
    cachedNum = cachedNum / max;
    cached -= 1;
    return r;
*/
    //if (noRandom) return 1;
    //return randomInt(r, min, max);
}

breakpoint() __llvm__
{
    call void asm sideeffect "int $$3", ""()
    ret i32 0
}

tryStep(s: System, f)
{

//if (true) return false;

    var r1 = fastRandomInt(s.r^, 0, s.n-1);
    var r2 = fastRandomInt(s.r^, 0, s.n-1);
  
    //var r1 = 1;
    //var r2 = 0;
  /*
    if (noRandom)
    {
        r1 = 0;
        r2 = 0;
    }*/
    
    //var oldE = s.E;
    //s.E = computeEnergy(s);
    
    //breakpoint();

    //if (true) return true;
/*    
    var newE = energyWithSwappedRows(s, r1, r2);

    if (f(newE))
    {
        s.E = newE;
        diagUpdateRemoveRow(s, r1);
        diagUpdateRemoveRow(s, r2);
        swap(s.board[r1], s.board[r2]);
        //for (x in [r1, r2]) energyUpdateAddRow(s, x);
        diagUpdateAddRow(s, r1);
        diagUpdateAddRow(s, r2);
//if(s.E != computeEnergy(s)) println("ouch !");
        //testDiags(s);
        return true;
    }
//if(s.E != computeEnergy(s)) println("ouch !");
//if(s.E != computeEnergy(s)) println("ouch !");
        //testDiags(s);
    return false;

//    */

    //for (x in [r1, r2]) energyUpdateRemoveRow(s, x);
    energyUpdateRemoveRow(s, r1);
    energyUpdateRemoveRow(s, r2);
    swap(s.board[r1], s.board[r2]);
    //for (x in [r1, r2]) energyUpdateAddRow(s, x);
    energyUpdateAddRow(s, r1);
    energyUpdateAddRow(s, r2);

    //if(s.E != computeEnergy(s)) println("ouch !");


    if (not f(s.E))
    //if (true)
    {
        // rejected: undo step
        //for (x in [r1, r2]) energyUpdateRemoveRow(s, x);
        energyUpdateRemoveRow(s, r1);
        energyUpdateRemoveRow(s, r2);
        swap(s.board[r1], s.board[r2]);
        //for (x in [r1, r2]) energyUpdateAddRow(s, x);
        energyUpdateAddRow(s, r1);
        energyUpdateAddRow(s, r2);

        //if(s.E != computeEnergy(s)) println("ouch2 !");
        
        //s.E = oldE;
        return false;
    }

    return true;
    // */
}

energyWithSwappedRows(s: System, r1, r2)
{

    var collision = false;

    var a = s.board[r1]-r1+s.n;
    var b = s.board[r2]-r2+s.n;
    var c = s.board[r2]-r1+s.n;
    var d = s.board[r1]-r2+s.n;

    //println([a,b,c,d]);

    if (a==b or a==c or a==d or b==c or b==d or c==d)
        collision = true;

    var a = s.board[r1]-r1+s.n;
    var b = s.board[r2]-r2+s.n;
    var c = s.board[r2]-r1+s.n;
    var d = s.board[r1]-r2+s.n;

    //println([a,b,c,d]);
    
    if (a==b or a==c or a==d or b==c or b==d or c==d)
        collision = true;

    //collision = false;

    if (collision)
    {
/*
    swap(s.board[r1], s.board[r2]);
    var newE = computeEnergy(s);
    swap(s.board[r1], s.board[r2]);
  */
  
    //var olds = s;
  
    energyUpdateRemoveRow(s, r1);
    energyUpdateRemoveRow(s, r2);
    swap(s.board[r1], s.board[r2]);
    energyUpdateAddRow(s, r1);
    energyUpdateAddRow(s, r2);

    var newE = s.E;
        
    energyUpdateRemoveRow(s, r1);
    energyUpdateRemoveRow(s, r2);
    swap(s.board[r1], s.board[r2]);
    energyUpdateAddRow(s, r1);
    energyUpdateAddRow(s, r2);
        
    //s = olds;
    
    return newE;    
    }

    //println("test");
    //testDiags(s);

    var newE = s.E;
    newE -= if (s.rdiag[s.board[r1]-r1+s.n] > 1) 1 else 0;
    newE -= if (s.ldiag[s.board[r1]+r1] > 1) 1 else 0;
    newE -= if (s.rdiag[s.board[r2]-r2+s.n] > 1) 1 else 0;
    newE -= if (s.ldiag[s.board[r2]+r2] > 1) 1 else 0;
/*
    // diags
    newE += if ((s.board[r2]-r2 == s.board[r1]-r1) s.rdiag[s.board[r2]-r2+s.n] == 2) 1 else 0;
    newE += if ((s.board[r2]+r2 == s.board[r1]+r1) s.ldiag[s.board[r2]+r2] == 2) 1 else 0;
*/

    //println(newE);

    newE += if (s.rdiag[s.board[r2]-r1+s.n] >= 1) 1 else 0;
    newE += if (s.ldiag[s.board[r2]+r1] >= 1) 1 else 0;
    newE += if (s.rdiag[s.board[r1]-r2+s.n] >= 1) 1 else 0;
    newE += if (s.ldiag[s.board[r1]+r2] >= 1) 1 else 0;
/*
    // diags
    newE -= if ((s.board[r2]-r2 == s.board[r1]-r1) s.rdiag[s.board[r2]-r2+s.n] == 2) 1 else 0;
    newE -= if ((s.board[r2]+r2 == s.board[r1]+r1) s.ldiag[s.board[r2]+r2] == 2) 1 else 0;
*/
/*
    swap(s.board[r1], s.board[r2]);
    var newE2 = computeEnergy(s);
    swap(s.board[r1], s.board[r2]);

    if (newE != newE2)
    {
        println([newE, newE2]);
        println([r1, r2]);
        print(s);
    swap(s.board[r1], s.board[r2]);
        print(s);
    swap(s.board[r1], s.board[r2]);
        breakpoint();
    }
*/
    return newE;

    // implement...

    //if (s.rdiag[s.board[x]-x+s.n] > 1) s.E -= 1;
    //s.E -= if (s.rdiag[s.board[x]-x+s.n] > 1) 1 else 0;
    //s.rdiag[s.board[x]-x+s.n] -= 1;
    
    //if (s.ldiag[s.board[x]+x] > 1) s.E -= 1;
    //s.E -= if (s.ldiag[s.board[x]+x] > 1) 1 else 0;
    //s.ldiag[s.board[x]+x] -= 1;
}

diagUpdateRemoveRow(s: System, x)
{
    s.rdiag[s.board[x]-x+s.n] -= 1;
    s.ldiag[s.board[x]+x] -= 1;
}

diagUpdateAddRow(s: System, x)
{
    s.rdiag[s.board[x]-x+s.n] += 1;
    s.ldiag[s.board[x]+x] += 1;
}

energyUpdateRemoveRow(s: System, x)
{
    //if (s.rdiag[s.board[x]-x+s.n] > 1) s.E -= 1;
    s.E -= if (s.rdiag[s.board[x]-x+s.n] > 1) 1 else 0;
    s.rdiag[s.board[x]-x+s.n] -= 1;
    
    //if (s.ldiag[s.board[x]+x] > 1) s.E -= 1;
    s.E -= if (s.ldiag[s.board[x]+x] > 1) 1 else 0;
    s.ldiag[s.board[x]+x] -= 1;
}

energyUpdateAddRow(s: System, x)
{
    s.rdiag[s.board[x]-x+s.n] += 1;
    //if (s.rdiag[s.board[x]-x+s.n] > 1) s.E += 1;
    s.E += if (s.rdiag[s.board[x]-x+s.n] > 1) 1 else 0;

    s.ldiag[s.board[x]+x] += 1;
    //if (s.ldiag[s.board[x]+x] > 1) s.E += 1;
    s.E += if (s.ldiag[s.board[x]+x] > 1) 1 else 0;
}

testDiags(s: System)
{
    var oldr = s.rdiag;
    var oldl = s.ldiag;

    for(v in s.rdiag) v = 0;
    for(v in s.ldiag) v = 0;
    
    for(x in range(s.n))
    {
        s.rdiag[s.board[x]-x+s.n] += 1;
        s.ldiag[s.board[x]+x] += 1;
    }

    for (x in range(size(oldr))) if (oldr[x] != s.rdiag[x]) println("rdiag corrupted");
    for (x in range(size(oldl))) if (oldl[x] != s.ldiag[x]) println("ldiag corrupted");

    // avoid side effect
    s.rdiag = oldr;
    s.ldiag = oldl;
}

setDiags(s: System)
{
    for(v in s.rdiag) v = 0;
    for(v in s.ldiag) v = 0;
    
    for(x in range(s.n))
    {
        s.rdiag[s.board[x]-x+s.n] += 1;
        s.ldiag[s.board[x]+x] += 1;
    }
}

computeEnergy(s: System)
{
    var oldr = s.rdiag;
    var oldl = s.ldiag;

    for(v in s.rdiag) v = 0;
    for(v in s.ldiag) v = 0;
    
    for(x in range(s.n))
    {
        s.rdiag[s.board[x]-x+s.n] += 1;
        s.ldiag[s.board[x]+x] += 1;
    }

    //for (x in range(size(oldr))) if (oldr[x] != s.rdiag[x]) println("rdiag corrupted");
    //for (x in range(size(oldl))) if (oldl[x] != s.ldiag[x]) println("ldiag corrupted");
    
    var E = 0;
    for (v in s.rdiag) E += max(v-1,0);
    for (v in s.ldiag) E += max(v-1,0);

    // avoid side effect
    s.rdiag = oldr;
    s.ldiag = oldl;

    return E;
}

overload print(s: System)
{
    println("board:");
    for (v in s.board) 
    {
        //println(v);
        for (x in range(v)) print(". ");
        print("x ");
        for (x in range(s.n-v-1)) print(". ");
        println();
    }
    //println("energy: ", s.E, "[", computeEnergy(s), "]");
    println("energy: ", computeEnergy(s));
}