package puzzle;

import lombok.AllArgsConstructor;
import lombok.val;
import puzzle.Export.Clued;
import puzzle.Export.Gridded;

import java.sql.Array;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Objects;
import java.util.stream.IntStream;

import static java.lang.Long.*;
import static puzzle.SwedishGenerator.*;

public record Masker(Rng rng, int[] stack, Clues cache) {
   
   public static final int[][] MUTATE_RI = new int[SwedishGenerator.SIZE][625];
   
   static {
      for (int i = 0; i < SwedishGenerator.SIZE; i++) {
         int k = 0;
         for (int dr1 = -2; dr1 <= 2; dr1++)
            for (int dr2 = -2; dr2 <= 2; dr2++)
               for (int dc1 = -2; dc1 <= 2; dc1++)
                  for (int dc2 = -2; dc2 <= 2; dc2++) {
                     val ti = IT[i];
                     MUTATE_RI[i][k++] = Grid.offset(SwedishGenerator.clamp(ti.r() + dr1 + dr2, 0, R - 1),
                                                     SwedishGenerator.clamp(ti.c() + dc1 + dc2, 0, C - 1));
                  }
      }
   }
   // slice ray to stop before first clue, depending on direction monotonicity
   // right/down => increasing indices; up/left => decreasing indices
   // first clue is highest index among hits (hi first, then lo)
   private static void processSlotRev(Clues c, SlotVisitor visitor, int key) {
      long rayLo = PATH_LO[key];
      long rayHi = PATH_HI[key];
      // only consider clue cells
      long hitsLo = rayLo & c.lo;
      long hitsHi = rayHi & c.hi;
      
      if (hitsHi != X) {
         int  msb  = 63 - numberOfLeadingZeros(hitsHi);
         long stop = 1L << msb;
         rayHi &= -(stop << 1); // keep bits > stop
         rayLo = 0;                   // lo indices are below stop
      } else if (hitsLo != X) {
         int  msb  = 63 - numberOfLeadingZeros(hitsLo);
         long stop = 1L << msb;
         rayLo &= -(stop << 1);
      }
    //  if (Long.bitCount(rayLo) + Long.bitCount(rayHi) > 1)
         visitor.visit(key, rayLo, rayHi);
   }
   private static void processSlot(Clues c, SlotVisitor visitor, int key) {
      long rayLo  = PATH_LO[key];
      long rayHi  = PATH_HI[key];
      long hitsLo = rayLo & c.lo;
      long hitsHi = rayHi & c.hi;
      
      if (hitsLo != X) {
         long stop = 1L << numberOfTrailingZeros(hitsLo);
         rayLo &= (stop - 1);
         rayHi = 0; // any hi is beyond the stop
      } else if (hitsHi != X) {
         long stop = 1L << numberOfTrailingZeros(hitsHi);
         // keep all lo (lo indices are < any hi index), but cut hi below stop
         rayHi &= (stop - 1);
      }
     // if (Long.bitCount(rayLo) + Long.bitCount(rayHi) > 1)
         visitor.visit(key, rayLo, rayHi);
   }
   public static Slot[] extractSlots(Clues grid, DictEntry[] index) {
      var   slots = new ArrayList<Slot>(grid.clueCount());
      grid.forEachSlot((key, lo, hi) -> slots.add(Slot.from(key, lo, hi, Objects.requireNonNull(index[Slot.length(lo, hi)]))));
      return slots.toArray(Slot[]::new);
   }
   public static Slotinfo[] slots(Clues mask, DictEntry[] index) {
      var slots = Masker.extractSlots(mask, index);
      return Masker.scoreSlots(slots);
   }
   public static Slotinfo[] scoreSlots(Slot[] slots) {
      val        count    = new byte[SwedishGenerator.SIZE];
      Slotinfo[] slotInfo = new Slotinfo[slots.length];
      for (var s : slots) {
         for (long b = s.lo; b != X; b &= b - 1) count[numberOfTrailingZeros(b)]++;
         for (long b = s.hi; b != X; b &= b - 1) count[64 | numberOfTrailingZeros(b)]++;
      }
      for (int i = 0; i < slots.length; i++) {
         var slot = slots[i];
         slotInfo[i] = new Slotinfo(slot.key, slot.lo, slot.hi, slotScore(count, slot.lo, slot.hi), new Assign(), slot.entry);
      }
      return slotInfo;
   }
   public static int slotScore(byte[] count, long lo, long hi) {
      int cross = 0;
      for (long b = lo; b != X; b &= b - 1) cross += (count[numberOfTrailingZeros(b)] - 1);
      for (long b = hi; b != X; b &= b - 1) cross += (count[64 | numberOfTrailingZeros(b)] - 1);
      return cross * 10 + Slot.length(lo, hi);
   }
   
   public long maskFitness(final Clues grid, int clueSize) {
      
      long    cHLo     = 0L, cHHi = 0L, cVLo = 0L, cVHi = 0L;
      long    cHLo2    = 0L, cHHi2 = 0L, cVLo2 = 0L, cVHi2 = 0L;
      long    lo_cl    = grid.lo, hi_cl = grid.hi;
      long    penalty  = (((long) Math.abs(grid.clueCount() - clueSize)) * 16000L);
      boolean hasSlots = false;
      
      for (long bits = lo_cl; bits != X; bits &= bits - 1) {
         long lsb     = bits & -bits;
         int  clueIdx = numberOfTrailingZeros(lsb);
         int  v       = (grid.vlo & lsb) != 0 ? 1 : 0;
         int  r       = (grid.rlo & lsb) != 0 ? 1 : 0;
         int  key     = Slot.packSlotKey(clueIdx, (r << 1) | v);
         long rLo     = PATH_LO[key], rHi = PATH_HI[key];
         long hLo     = rLo & lo_cl, hHi = rHi & hi_cl;
         if (Slotinfo.increasing(key)) {
            if (hLo != X) {
               rLo &= ((1L << numberOfTrailingZeros(hLo)) - 1);
               rHi = 0;
            } else if (hHi != X) rHi &= ((1L << numberOfTrailingZeros(hHi)) - 1);
         } else if (hHi != X) {
            int msb = 63 - numberOfLeadingZeros(hHi);
            rHi &= -(1L << msb << 1);
            rLo = 0;
         } else if (hLo != X) {
            int msb = 63 - numberOfLeadingZeros(hLo);
            rLo &= -(1L << msb << 1);
         }
         if ((rLo | rHi) != X) {
            hasSlots = true;
            if (Slot.horiz(key)) {
               cHLo2 |= (cHLo & rLo); cHHi2 |= (cHHi & rHi);
               cHLo |= rLo;
               cHHi |= rHi;
            } else {
               cVLo2 |= (cVLo & rLo); cVHi2 |= (cVHi & rHi);
               cVLo |= rLo;
               cVHi |= rHi;
            }
            if ((bitCount(rLo) + bitCount(rHi)) < MIN_LEN) penalty += 8000;
         } else penalty += 25000;
      }
      for (long bits = hi_cl; bits != X; bits &= bits - 1) {
         long lsb     = bits & -bits;
         int  clueIdx = numberOfTrailingZeros(lsb);
         int  v       = (grid.vhi & lsb) != 0 ? 1 : 0;
         int  r       = (grid.rhi & lsb) != 0 ? 1 : 0;
         int  key     = Slot.packSlotKey(64 | clueIdx, (r << 1) | v);
         long rLo     = PATH_LO[key], rHi = PATH_HI[key];
         long hLo     = rLo & lo_cl, hHi = rHi & hi_cl;
         if (Slotinfo.increasing(key)) {
            if (hLo != X) {
               rLo &= ((1L << numberOfTrailingZeros(hLo)) - 1);
               rHi = 0;
            } else if (hHi != X) rHi &= ((1L << numberOfTrailingZeros(hHi)) - 1);
         } else if (hHi != X) {
            int msb = 63 - numberOfLeadingZeros(hHi);
            rHi &= -(1L << msb << 1);
            rLo = 0;
         } else if (hLo != X) {
            int msb = 63 - numberOfLeadingZeros(hLo);
            rLo &= -(1L << msb << 1);
         }
         if ((rLo | rHi) != X) {
            hasSlots = true;
            if (Slot.horiz(key)) {
               cHLo2 |= (cHLo & rLo); cHHi2 |= (cHHi & rHi);
               cHLo |= rLo;
               cHHi |= rHi;
            } else {
               cVLo2 |= (cVLo & rLo); cVHi2 |= (cVHi & rHi);
               cVLo |= rLo;
               cVHi |= rHi;
            }
            if ((bitCount(rLo) + bitCount(rHi)) < MIN_LEN) penalty += 8000;
         } else penalty += 25000;
      }
      
      if (!hasSlots) return 1_000_000_000L;
      long seenLo = X, seenHi = X;

// loop over beide helften
      for (int base = 0, size, sp, cur; base <= 64; base += 64) {
         long clueMask = (base == 0) ? lo_cl : hi_cl;
         long seenMask = (base == 0) ? seenLo : seenHi;
         
         // "unseen clues" in deze helft
         for (long bits = clueMask & ~seenMask, nLo, nHi; bits != X; bits &= bits - 1) {
            int clueIdx = base | numberOfTrailingZeros(bits);
            
            // start nieuwe component
            size     = 0;
            stack[0] = clueIdx;
            sp       = 1;
            
            // mark seen
            if ((clueIdx & 64) == 0) seenLo |= 1L << clueIdx;
            else seenHi |= 1L << (clueIdx & 63);
            
            // flood fill / bfs
            while (sp > 0) {
               cur = stack[--sp];
               size++;
               
               // neighbors als 2x long masks
               nLo = NBR8_PACKED_LO[cur];
               nHi = NBR8_PACKED_HI[cur];
               
               // filter: alleen clues, en nog niet seen
               nLo &= lo_cl & ~seenLo;
               nHi &= hi_cl & ~seenHi;
               
               // push lo-neighbors
               while (nLo != X) {
                  long lsb  = nLo & -nLo;
                  int  nidx = numberOfTrailingZeros(nLo); // 0..63
                  seenLo |= lsb;
                  
                  stack[sp++] = nidx;
                  
                  nLo &= nLo - 1;
               }
               
               // push hi-neighbors
               while (nHi != X) {
                  long lsb  = nHi & -nHi;
                  int  nidx = 64 | numberOfTrailingZeros(nHi); // 64..127
                  seenHi |= lsb;
                  
                  stack[sp++] = nidx;
                  
                  nHi &= nHi - 1;
               }
            }
            
            if (size >= 2) penalty += (size - 1L) * 120L;
         }
      }
      
      for (long bits = ~lo_cl & MASK_LO; bits != X; bits &= bits - 1) {
         int clueIdx = numberOfTrailingZeros(bits);
         var rci     = IT[clueIdx];
         if ((4 - rci.nbrCount()) + bitCount(rci.n1() & lo_cl) + bitCount(rci.n2() & hi_cl) >= 3) penalty += 400;
         boolean h = (cHLo & (1L << clueIdx)) != X;
         boolean v = (cVLo & (1L << clueIdx)) != X;
         if (!h && !v) penalty += 1500;
         else if (h && v) { /* ok */ }
         else if (((h ? cHLo2 : cVLo2) & (1L << clueIdx)) != X) penalty += 600;
         else penalty += 200;
      }
      for (long bits = ~hi_cl & MASK_HI; bits != X; bits &= bits - 1) {
         int clueIdx = numberOfTrailingZeros(bits);
         var rci     = IT[64 | clueIdx];
         if ((4 - rci.nbrCount()) + bitCount(rci.n1() & lo_cl) + bitCount(rci.n2() & hi_cl) >= 3) penalty += 400;
         boolean h = (cHHi & (1L << clueIdx)) != X;
         boolean v = (cVHi & (1L << clueIdx)) != X;
         if (!h && !v) penalty += 1500;
         else if (h && v) { /* ok */ }
         else if (((h ? cHHi2 : cVHi2) & (1L << clueIdx)) != X) penalty += 600;
         else penalty += 200;
      }
      
      return penalty;
   }
   
   public Clues randomMask(final int clueSize) {
      var g = Clues.createEmpty();
      for (int placed = 0, guard = 0, ri; placed < clueSize && guard < 4000; guard++) {
         
         ri = rng.randint0_SIZE();
         if (isLo(ri)) {
            if (g.isClueLo(ri)) continue;
            var d_idx = rng.randint2bitByte();
            if (g.hasRoomForClue(OFFSETS_D_IDX[Slot.packSlotKey(ri, d_idx)])) {
               g.setClueLo(1L << ri, d_idx);
               placed++;
            }
         } else {
            if (g.isClueHi(ri)) continue;
            var d_idx = rng.randint2bitByte();
            if (g.hasRoomForClue(OFFSETS_D_IDX[Slot.packSlotKey(ri, d_idx)])) {
               g.setClueHi(1L << (ri & 63), d_idx);
               placed++;
            }
         }
         
      }
      return g;
   }
   
   public Clues mutate(Clues c) {
      var bytes = MUTATE_RI[rng.randint0_SIZE()];
      for (int k = 0, ri; k < 4; k++) {
         ri = bytes[rng.randint0_624()];
         if (isLo(ri)) {
            if (!c.cluelessLo(ri)) {
               var d_idx = rng.randint2bitByte();
               if (c.hasRoomForClue(OFFSETS_D_IDX[Slot.packSlotKey(ri, d_idx)])) c.setClueLo(1L << ri, d_idx);
            }
         } else if (!c.cluelessHi(ri)) {
            var d_idx = rng.randint2bitByte();
            if (c.hasRoomForClue(OFFSETS_D_IDX[Slot.packSlotKey(ri, d_idx)])) c.setClueHi(1L << (ri & 63), d_idx);
         }
         
      }
      return c;
   }
   
   public Clues crossover(Clues a, Clues other) {
      
      var theta = rng.nextFloat() * Math.PI;
      var nc    = Math.cos(theta);
      var nr    = Math.sin(theta);
      
      long maskLo = 0, maskHi = 0;
      for (var rci : IT)
         if ((rci.cross_r()) * nc + (rci.cross_c()) * nr < 0) {
            int i = rci.i();
            if ((i & 64) == 0) maskLo |= (1L << i);
            else maskHi |= (1L << (i - 64));
         }
      var c = new Clues(
              (a.lo & ~maskLo) | (other.lo & maskLo),
              (a.hi & ~maskHi) | (other.hi & maskHi),
              (a.vlo & ~maskLo) | (other.vlo & maskLo),
              (a.vhi & ~maskHi) | (other.vhi & maskHi),
              (a.rlo & ~maskLo) | (other.rlo & maskLo),
              (a.rhi & ~maskHi) | (other.rhi & maskHi));
      
      for (var l = c.lo & ~c.rlo & ~c.vlo; l != X; l &= l - 1) clearCluesLo(c, numberOfTrailingZeros(l), 0);
      for (var l = c.lo & ~c.rlo & c.vlo; l != X; l &= l - 1) clearCluesLo(c, numberOfTrailingZeros(l), 1);
      for (var l = c.lo & c.rlo & ~c.vlo; l != X; l &= l - 1) clearCluesLo(c, numberOfTrailingZeros(l), 2);
      for (var l = c.lo & c.rlo & c.vlo; l != X; l &= l - 1) clearCluesLo(c, numberOfTrailingZeros(l), 3);
      for (var h = c.hi & ~c.rhi & ~c.vhi; h != X; h &= h - 1) clearCluesHi(c, numberOfTrailingZeros(h), 0);
      for (var h = c.hi & ~c.rhi & c.vhi; h != X; h &= h - 1) clearCluesHi(c, numberOfTrailingZeros(h), 1);
      for (var h = c.hi & c.rhi & ~c.vhi; h != X; h &= h - 1) clearCluesHi(c, (numberOfTrailingZeros(h)), 2);
      for (var h = c.hi & c.rhi & c.vhi; h != X; h &= h - 1) clearCluesHi(c, (numberOfTrailingZeros(h)), 3);
      return c;
   }
   
   public static void clearCluesLo(Clues out, int idx, int d) { if (!out.hasRoomForClue(OFFSETS_D_IDX[Slot.packSlotKey(idx, d)])) out.clearClueLo(~(1L << idx)); }
   
   public static void clearCluesHi(Clues out, int idx, int d) { if (!out.hasRoomForClue(OFFSETS_D_IDX[Slot.packSlotKey(64 | idx, d)])) out.clearClueHi(~(1L << idx)); }
   
   public Clues hillclimb(Clues start, int clue_size, int limit) {
      var best  = start;
      var bestF = maskFitness(best, clue_size);
      var fails = 0;
      
      while (fails < limit) {
         cache.from(best);
         var cand = mutate(best);
         var f    = maskFitness(cand, clue_size);
         if (f < bestF) {
            best  = cand;
            bestF = f;
            fails = 0;
         } else {
            best.from(cache);
            fails++;
         }
      }
      return best;
   }
   
   public Clues generateMask(int clueSize, int popSize, int gens, int offspring) {
      class GridAndFit {
         
         Clues grid;
         long  fite = -1;
         GridAndFit(Clues grid) { this.grid = grid; }
         long fit() {
            if (fite == -1) this.fite = maskFitness(grid, clueSize);
            return this.fite;
         }
      }
      if (Main.VERBOSE) System.out.println("generateMask init pop: " + popSize + " clueSize: " + clueSize);
      var pop = new ArrayList<GridAndFit>();
      for (var i = 0; i < popSize; i++) {
         if (Thread.currentThread().isInterrupted()) return null;
         pop.add(new GridAndFit(hillclimb(randomMask(clueSize), clueSize, 180)));
      }
      
      for (var gen = 0; gen < gens; gen++) {
         if (Thread.currentThread().isInterrupted()) break;
         var children = new ArrayList<GridAndFit>();
         
         for (var k = 0; k < offspring; k++) {
            if (Thread.currentThread().isInterrupted()) break;
            var p1    = rng.rand(pop);
            var p2    = rng.rand(pop);
            var child = crossover(p1.grid, p2.grid);
            children.add(new GridAndFit(hillclimb(child, clueSize, 70)));
         }
         
         pop.addAll(children);
         pop.sort(Comparator.comparingLong(GridAndFit::fit));
         
         var next = new ArrayList<GridAndFit>();
         for (var cand : pop) {
            if (next.size() >= offspring) break;
            var ok = true;
            for (var kept : next)
               if (cand.grid.similarity(kept.grid) > 0.92) {
                  ok = false;
                  break;
               }
            if (ok) next.add(cand);
         }
         pop = next;
         
         if (Main.VERBOSE && (gen & 15) == 15) System.out.println("  gen " + gen + "/" + gens + " bestFitness=" + pop.get(0).fit());
      }
      if (pop.isEmpty()) return null;
      GridAndFit best = pop.get(0);
      for (int i = 1; i < pop.size(); i++) {
         var x = pop.get(i);
         if (x.fit() < best.fit()) best = x;
      }
      return best.grid;
   }//@formatter:off
   @FunctionalInterface public interface SlotVisitor { void visit(int key, long lo, long hi); }
   @AllArgsConstructor
   public static class Clues {
      
      long lo, hi, vlo, vhi, rlo, rhi;
      
      public static Clues createEmpty() { return new Clues(0, 0, 0, 0, 0, 0); }
      public static Clued parse(String s) {
         var c     = createEmpty();
         var lines = s.split("\n");
         for (int r = 0; r < Math.min(lines.length, R); r++) {
            var line = lines[r];
            for (int col = 0; col < Math.min(line.length(), C); col++) {
               char ch = line.charAt(col);
               if (ch >= '0' && ch <= '3') {
                  int  idx = Grid.offset(r, col);
                  byte dir = (byte) (ch - '0');
                  if ((idx & 64) == 0) c.setClueLo(1L << idx, dir);
                  else c.setClueHi(1L << (idx & 63), dir);
               }
            }
         }
         return new Clued(c);
      }
      public boolean cluelessLo(int idx) {
         if (!isClueLo(idx)) return false;
         clearClueLo(~(1L << idx));
         return true;
      }
      public boolean cluelessHi(int idx) {
         if (!isClueHi(idx)) return false;
         clearClueHi(~(1L << (idx & 63)));
         return true;
      }
      public boolean hasRoomForClue(long packed) { return (packed) != X && notClue(packed & 0x7FL) && notClue((packed >>> 7) & 0x7FL); }
      
      public void setClueLo(long mask, byte idx) {
         lo |= mask;
         if ((idx & 1) != 0) vlo |= mask;
         else vlo &= ~mask;
         if ((idx & 2) != 0) rlo |= mask;
         else rlo &= ~mask;
      }
      public void setClueHi(long mask, byte idx) {
         hi |= mask;
         if ((idx & 1) != 0) vhi |= mask;
         else vhi &= ~mask;
         if ((idx & 2) != 0) rhi |= mask;
         else rhi &= ~mask;
      }
      public void clearClueLo(long mask) {
         lo &= mask;
         vlo &= mask;
         rlo &= mask;
      }
      public void clearClueHi(long mask) {
         hi &= mask;
         vhi &= mask;
         rhi &= mask;
      }
      public boolean isClueLo(int index) { return ((lo >>> index) & 1L) != X; }
      public boolean isClueHi(int index) { return ((hi >>> (index & 63)) & 1L) != X; }
      public boolean notClue(long index) { return ((index & 64) == 0) ? ((lo >>> index) & 1L) == X : ((hi >>> (index & 63)) & 1L) == X; }
      public boolean notClue(int index)  { return ((index & 64) == 0) ? ((lo >>> index) & 1L) == X : ((hi >>> (index & 63)) & 1L) == X; }
      
      public int clueCount()             { return bitCount(lo) + bitCount(hi); }
      public double similarity(Clues b) {
         long matchLo = (~(lo ^ b.lo)) & (~lo | (~(vlo ^ b.vlo) & ~(rlo ^ b.rlo)));
         long matchHi = (~(hi ^ b.hi)) & (~hi | (~(vhi ^ b.vhi) & ~(rhi ^ b.rhi)));
         
         return (bitCount(matchLo & MASK_LO) + bitCount(matchHi & MASK_HI)) / SIZED;
      }
      public Grid toGrid() { return new Grid(new byte[SwedishGenerator.SIZE], lo, hi); }
      public void forEachSlot(SlotVisitor visitor) {
         for (var l = lo & ~rlo & vlo; l != X; l &= l - 1) processSlot(this, visitor, Slot.packSlotKey(numberOfTrailingZeros(l), 1));
         for (var l = lo & ~rlo & ~vlo; l != X; l &= l - 1) processSlot(this, visitor, Slot.packSlotKey(numberOfTrailingZeros(l), 0));
         for (var l = lo & rlo & ~vlo; l != X; l &= l - 1) processSlotRev(this, visitor, Slot.packSlotKey(numberOfTrailingZeros(l), 2));
         for (var l = lo & rlo & vlo; l != X; l &= l - 1) processSlotRev(this, visitor, Slot.packSlotKey(numberOfTrailingZeros(l), 3));
         for (var h = hi & ~rhi & vhi; h != X; h &= h - 1) processSlot(this, visitor, Slot.packSlotKey(64 | numberOfTrailingZeros(h), 1));
         for (var h = hi & ~rhi & ~vhi; h != X; h &= h - 1) processSlot(this, visitor, Slot.packSlotKey(64 | numberOfTrailingZeros(h), 0));
         for (var h = hi & rhi & ~vhi; h != X; h &= h - 1) processSlotRev(this, visitor, Slot.packSlotKey((64 | numberOfTrailingZeros(h)), 2));
         for (var h = hi & rhi & vhi; h != X; h &= h - 1) processSlotRev(this, visitor, Slot.packSlotKey((64 | numberOfTrailingZeros(h)), 3));
      }
      public Clues from(Clues best) {
         lo  = best.lo;
         hi  = best.hi;
         vlo = best.vlo;
         vhi = best.vhi;
         rlo = best.rlo;
         rhi = best.rhi;
         return this;
      }
   }
   
   public record Slot(int key, long lo, long hi, DictEntry entry) {
      
      static final int BIT_FOR_DIR = 2;
      static Slot from(int key, long lo, long hi, DictEntry entry) { return new Slot(key, lo, hi, entry); }
      public static int length(long lo, long hi)                   { return bitCount(lo) + bitCount(hi); }
      public static int clueIndex(int key)                         { return key >>> BIT_FOR_DIR; }
      public static int dir(int key)                               { return key & 3; }
      public IntStream walk()                                      { return Gridded.walk((byte) key, lo, hi); }
      public static boolean horiz(int d)                           { return (d & 1) != 0; }
      public static int packSlotKey(int idx, int d)                { return (idx << BIT_FOR_DIR) | d; }
   }
}