module NonogramSolverPar (solveSequential, solveParallelComputeBlocks, solveParallelGenerateBlocks, solveParallelComputeGenerate, solveParallelIterativeSolve, solveFullyParallel) where
    
import Data.Array
import Data.List (group, transpose, foldl')
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Bits ((.&.), (.|.))
import Parser (parseNonogram)
import NonogramTypes

import Control.Parallel.Strategies (parMap, rdeepseq, rseq)

computeBlocksSeq :: Int -> [Int] -> [[Int]]
computeBlocksSeq lineLength blockLengths = placeBlocks blockLengths 0
  where
    -- place blocks
    placeBlocks :: [Int] -> Int -> [[Int]]
    placeBlocks [] _ = [[]] -- No blocks left to place
    placeBlocks (b:bs) start
      | start + remainingLength > lineLength = []
      | otherwise = do
          pos <- [start .. lineLength - remainingLength]
          rest <- placeBlocks bs (pos + b + 1) 
          return (pos : rest)
      where
        -- remaining length
        remainingLength = sum (b : bs) + length bs

-- Compute all valid placements of blocks w/ chunked parallelization
computeBlocksPar :: Int -> [Int] -> [[Int]]
computeBlocksPar lineLength blockLengths =
    concat $ parMap rdeepseq processChunk (chunkList chunkSize [0 .. lineLength - totalRemainingLength])
  where
    totalRemainingLength = sum blockLengths + length blockLengths - 1
    chunkSize = 10 

    -- Divide list into chunks
    chunkList :: Int -> [a] -> [[a]]
    chunkList _ [] = []
    chunkList n xs = take n xs : chunkList n (drop n xs)

    -- Process chunk of starting positions
    processChunk :: [Int] -> [[Int]]
    processChunk positions = concatMap processPosition positions

    -- Process single starting position
    processPosition :: Int -> [[Int]]
    processPosition start = placeBlocks blockLengths start

    --Helper to place blocks
    placeBlocks :: [Int] -> Int -> [[Int]]
    placeBlocks [] _ = [[]] -- No blocks left to place
    placeBlocks (b:bs) start
      | start + remainingLength > lineLength = []
      | otherwise = do
          pos <- [start .. lineLength - remainingLength]
          rest <- placeBlocks bs (pos + b + 1) 
          return (pos : rest)
      where
        remainingLength = sum (b : bs) + length bs 



generateBlocksSeq :: [[Int]] -> [Int] -> Int -> [[Int]]
generateBlocksSeq blockStarts blockSizes totalLength =
    map (generateBinaryArray blockSizes totalLength) blockStarts
  where
    generateBinaryArray :: [Int] -> Int -> [Int] -> [Int]
    generateBinaryArray sizes len starts = foldl placeBlock (replicate len 0) (zip starts sizes)

    placeBlock :: [Int] -> (Int, Int) -> [Int]
    placeBlock arr (start, size) =
        take start arr ++ replicate size 1 ++ drop (start + size) arr

-- Parallelized generateBlocks with chunking
generateBlocksPar :: [[Int]] -> [Int] -> Int -> [[Int]]
generateBlocksPar blockStarts blockSizes totalLength =
    concat $ parMap rdeepseq processChunk (chunkList chunkSize blockStarts)
  where
    chunkSize = 10 

    -- Divide list into chunks 
    chunkList :: Int -> [a] -> [[a]]
    chunkList _ [] = []
    chunkList n xs = take n xs : chunkList n (drop n xs)

    -- Process chunk of blockStarts
    processChunk :: [[Int]] -> [[Int]]
    processChunk chunk = map (generateBinaryArray blockSizes totalLength) chunk

    -- Generate binary array for single starting configuration
    generateBinaryArray :: [Int] -> Int -> [Int] -> [Int]
    generateBinaryArray sizes len starts =
        foldl' placeBlock (replicate len 0) (zip starts sizes)

    -- Place single block in  array
    placeBlock :: [Int] -> (Int, Int) -> [Int]
    placeBlock arr (start, size) =
        take start arr ++ replicate size 1 ++ drop (start + size) arr

        


validGroups :: [Int] -> String -> Bool
validGroups expected line = groupLengths == expected
  where
    -- Filter out zeroes and compute lengths of groups of ones
    groupLengths = map length . filter (all (== '1')) $ group line

--------------------------------------------------------------------------------------------------------------------------------------------------------

type Constraint = [Int]
type PartialSolution = Array (Int, Int) Int

-- Helper function: Extract groups of 1s from a row/column as strings
extractGroupsAsString :: [Int] -> String
extractGroupsAsString = map toChar
  where
    toChar 1 = '1'
    toChar 0 = '0'
    toChar _ = '-'

-- Function to validate single row or column
validateLine :: [Int] -> Constraint -> Bool
validateLine line constraint
  | all (== -1) line = True -- Skip untouched lines
  | sum (filter (== 1) line) > sum constraint = False -- Exceeding constraints
  | any (\x -> x == -1) line && sum (filter (== 1) line) <= sum constraint = True -- Partially solved line
  | otherwise = validGroups constraint (extractGroupsAsString line)

-- Main validation function
valid :: [Constraint] -> [Constraint] -> PartialSolution -> Set.Set Int -> Set.Set Int -> Bool
valid rowArgs colArgs partialSolution completedRows completedCols =
  let
    -- Get dimensions of grid
    ((r1, c1), (r2, c2)) = bounds partialSolution
    -- Extract localRows and columns from partial solution
    localRows = [[partialSolution ! (r, c) | c <- [c1..c2]] | r <- [r1..r2]]
    cols = transpose localRows
    -- Validate localRows
    validRows = all (\(r, row) ->
      let
        constraint = rowArgs !! r
      in
        if r `Set.member` completedRows
        then validGroups constraint (extractGroupsAsString row)
        else validateLine row constraint
      ) $ zip [0..] localRows
    -- Validate columns
    validCols = all (\(c, col) ->
      let
        constraint = colArgs !! c
      in
        if c `Set.member` completedCols
        then validGroups constraint (extractGroupsAsString col)
        else validateLine col constraint
      ) $ zip [0..] cols

  in
    validRows && validCols

--------------------------------------------------------------------------------------------------------------------------------------------------------

inferValuesSeq :: PartialSolution -> (Array Int (Set.Set [Int]), Array Int (Set.Set [Int])) -> PartialSolution
inferValuesSeq partialSolution (rowPlacements, colPlacements) =
    let ((r1, c1), (r2, c2)) = bounds partialSolution
        numCols = c2 - c1 + 1
        numRows = r2 - r1 + 1

        -- Update localRows
        updatedRows = partialSolution // concat (parMap rseq updateRow [r1..r2])
          where
            updateRow r =
                let placements = Set.toList $ rowPlacements ! r
                in if null placements
                   then []  -- If there are no placements, skip updates for this row
                   else
                       let rowOnes = foldl1 (zipWith (.&.)) placements
                           rowZeros = foldl1 (zipWith (.|.)) placements
                           inferredRow = inferRowOrCol rowOnes rowZeros
                       in [((r, c), inferredRow !! c) | c <- [0..numCols - 1], partialSolution ! (r, c) == -1]

        -- Update columns
        updatedCols = updatedRows // concatMap updateCol [c1..c2]
          where
            updateCol c =
                let placements = Set.toList $ colPlacements ! c
                in if null placements
                   then []  -- If there are no placements, skip updates for this column
                   else
                       let colOnes = foldl1 (zipWith (.&.)) placements
                           colZeros = foldl1 (zipWith (.|.)) placements
                           inferredCol = inferRowOrCol colOnes colZeros
                       in [((r, c), inferredCol !! r) | r <- [0..numRows - 1], updatedRows ! (r, c) == -1]

    in updatedCols


inferValuesPar :: PartialSolution -> (Array Int (Set.Set [Int]), Array Int (Set.Set [Int])) -> PartialSolution
inferValuesPar partialSolution (rowPlacements, colPlacements) =
    let ((r1, c1), (r2, c2)) = bounds partialSolution
        numCols = c2 - c1 + 1

        -- Parallel update for localRows
        updatedRows = partialSolution // concat (parMap rseq updateRow [r1..r2])
          where
            updateRow r =
                let placements = Set.toList $ rowPlacements ! r
                in if null placements
                   then []  -- If there are no placements, skip updates for this row
                   else
                       let rowOnes = foldl1 (zipWith (.&.)) placements
                           rowZeros = foldl1 (zipWith (.|.)) placements
                           inferredRow = inferRowOrCol rowOnes rowZeros
                       in [((r, c), inferredRow !! c) | c <- [0..numCols - 1], partialSolution ! (r, c) == -1]

        -- Sequential update for columns (after localRows are done)
        numRows = r2 - r1 + 1
        updatedCols = updatedRows // concatMap updateCol [c1..c2]
          where
            updateCol c =
                let placements = Set.toList $ colPlacements ! c
                in if null placements
                   then []  -- If there are no placements, skip updates for this column
                   else
                       let colOnes = foldl1 (zipWith (.&.)) placements
                           colZeros = foldl1 (zipWith (.|.)) placements
                           inferredCol = inferRowOrCol colOnes colZeros
                       in [((r, c), inferredCol !! r) | r <- [0..numRows - 1], updatedRows ! (r, c) == -1]

    in updatedCols

-- Infer result for row or column based on bitwise results
inferRowOrCol :: [Int] -> [Int] -> [Int]
inferRowOrCol ones zeros =
    zipWith resolveCell ones zeros
  where
    resolveCell 0 1 = -1   -- unknown
    resolveCell 1 1 = 1    -- filled (from ones)
    resolveCell 0 0 = 0    -- empty (from zeros)
    resolveCell _ _ = -1   -- unexpected
--------------------------------------------------------------------------------------------------------------------------------------------------------

-- Determine completed localRows and columns in solution array
-- Input: 2D array representing current grid state
-- Output: A tuple of two sets. localRows/colCompleted: indices of localRows/cols that are completed, 
-- How: Iterate through localRows and columns, check if all cells are not -1. 
updateCompletions :: Array (Int, Int) Int -> (Set Int, Set Int)
updateCompletions solutionArray =
    let ((r1, c1), (r2, c2)) = bounds solutionArray
        -- Extract localRows and check if all values in each row are not -1
        rowsCompleted = Set.fromList [r | r <- [r1..r2], all (/= -1) [solutionArray ! (r, c) | c <- [c1..c2]]]
        -- Extract columns and check if all values in each column are not -1
        colsCompleted = Set.fromList [c | c <- [c1..c2], all (/= -1) [solutionArray ! (r, c) | r <- [r1..r2]]]
    in (rowsCompleted, colsCompleted)


type Placement = [Int]
type PlacementsDict = (Array Int (Set.Set [Int]), Array Int (Set.Set [Int])) -- (rowPlacements, colPlacements)

-- Inputs: solutionArray, rowPlacements & colPlacements, completedRows & completedCols
-- Output: updated Bool (indicates whether it made any updates), PlacementsDict (updated rowPlacements & colPlacements)
-- Calls: updateRow, isValidRow, updateCol, isValidCol
updatePlacements :: Array (Int, Int) Int -> PlacementsDict -> Set.Set Int -> Set.Set Int -> (Bool, PlacementsDict)
updatePlacements solutionArray (rowPlacements, colPlacements) completedRows completedCols =
    let
        -- Update localRows
        (updatedRows, updatedRowPlacements) = processRows solutionArray rowPlacements completedRows
        -- Update columns
        (updatedCols, updatedColPlacements) = processColumns solutionArray colPlacements completedCols
        -- Combine results
        updated = updatedRows || updatedCols
    in
        (updated, (updatedRowPlacements, updatedColPlacements))

-- Process localRows: remove completed localRows and validate remaining placements
processRows :: Array (Int, Int) Int -> Array Int (Set.Set Placement) -> Set.Set Int -> (Bool, Array Int (Set.Set Placement))
processRows solutionArray rowPlacements completedRows =
    let
        indexesForDeletion = [i | i <- indices rowPlacements, i `Set.member` completedRows]
        validatePlacement rowIdx placement = all isValidCell (zip [0..] placement)
            where
                isValidCell (cellIdx, value) =
                    (value /= 0 || solutionArray ! (rowIdx, cellIdx) /= 1) &&
                    (value /= 1 || solutionArray ! (rowIdx, cellIdx) /= 0)
        processRow (rowUpdated, acc) rowIdx
            | rowIdx `Set.member` completedRows = (rowUpdated, acc // [(rowIdx, Set.empty)])
            | otherwise =
                let
                    validPlacements = Set.filter (validatePlacement rowIdx) (rowPlacements ! rowIdx)
                in
                    if validPlacements /= rowPlacements ! rowIdx
                        then (True, acc // [(rowIdx, validPlacements)]) -- make updates to acc array
                        else (rowUpdated, acc)
        (updated, newPlacements) = foldl' processRow (False, rowPlacements) (indices rowPlacements)
    in
        (updated, newPlacements)


-- Process columns: remove completed columns and validate remaining placements
processColumns :: Array (Int, Int) Int -> Array Int (Set.Set Placement) -> Set.Set Int -> (Bool, Array Int (Set.Set Placement))
processColumns solutionArray colPlacements completedCols =
    let
        indexesForDeletion = [i | i <- indices colPlacements, i `Set.member` completedCols] -- create a list of indices to remove later
        validatePlacement colIdx placement = all isValidCell (zip [0..] placement) -- colIdx is index of column being validated, placement is a candidate column placement
            where -- iterates over each cell in placement to get row indx and value
                isValidCell (cellIdx, value) =
                    (value /= 0 || solutionArray ! (cellIdx, colIdx) /= 1) &&
                    (value /= 1 || solutionArray ! (cellIdx, colIdx) /= 0)
        processCol (colUpdated, acc) colIdx
            | colIdx `Set.member` completedCols = (colUpdated, acc // [(colIdx, Set.empty)]) -- mark column as empty if completed
            | otherwise =
                let
                    validPlacements = Set.filter (validatePlacement colIdx) (colPlacements ! colIdx)
                in
                    if validPlacements /= colPlacements ! colIdx
                        then (True, acc // [(colIdx, validPlacements)])
                        else (colUpdated, acc)
        (updated, newPlacements) = foldl' processCol (False, colPlacements) (indices colPlacements)
    in
        (updated, newPlacements)



--------------------------------------------------------------------------------------------------------------------------------------------------------
-- Output: a list of PartialSolution(s)
backtrack :: PartialSolution
          -> Array Int (Set [Int])  -- Row placements
          -> [Constraint]           -- Row constraints
          -> [Constraint]           -- Column constraints
          -> Set Int                -- Completed localRows
          -> Set Int                -- Completed columns
          -> [PartialSolution]      -- Accumulated solutions
          -> [PartialSolution]      -- Final list of valid solutions
backtrack solutionArray rowPlacements rowArgs colArgs completedRows completedCols solutions
    | Set.size completedRows == length rowArgs =
        -- Base case: if all localRows are completed, add solution to list if unique
        if not (any (== solutionArray) solutions)
        then solutionArray : solutions
        else solutions
    | otherwise =
        -- Iterate over localRows in rowPlacements
        foldl tryPlacement solutions (assocs rowPlacements)
  where
    -- Try a placement for a given row
    tryPlacement :: [PartialSolution] -> (Int, Set.Set [Int]) -> [PartialSolution]
    tryPlacement solList (row, options) =
        if Set.member row completedRows
        then solList -- Skip completed localRows
        else foldl (tryOption row) solList (Set.toList options)

    -- Try an option for a specific row
    tryOption :: Int -> [PartialSolution] -> [Int] -> [PartialSolution]
    tryOption row solList option =
        let
            originalRow = [solutionArray ! (row, col) | col <- colIndices]
            updatedSolution = solutionArray // [((row, col), option !! (col - c1)) | col <- colIndices]
            (completedRowsNext, completedColsNext) = updateCompletions updatedSolution
            newRowPlacements = rowPlacements // [(row, Set.empty)]
        in
        if not (valid rowArgs colArgs updatedSolution completedRowsNext completedColsNext)
        then solList
        else
            -- Append solutions returned from recursive call to current list
            backtrack updatedSolution newRowPlacements rowArgs colArgs completedRowsNext completedColsNext solList
        where
            ((_, c1), (_, c2)) = bounds solutionArray
            colIndices = [c1..c2]


iterativeSolveSeq :: PartialSolution
               -> PlacementsDict       -- Row and column placements
               -> [Constraint]         -- Row constraints
               -> [Constraint]         -- Column constraints
               -> Set Int              -- Completed localRows
               -> Set Int              -- Completed columns
               -> (PartialSolution, PlacementsDict, Set Int, Set Int)
iterativeSolveSeq solutionArray (rowPlacements, colPlacements) rowArgs colArgs completedRows completedCols =
    let
        -- Infer values and update completions
        inferredSolution = inferValuesSeq solutionArray (rowPlacements, colPlacements)

        -- Update completions
        (newCompletedRows, newCompletedCols) = updateCompletions inferredSolution

        -- Update placements
        (updatedFlag, newPlacements) = updatePlacements inferredSolution (rowPlacements, colPlacements) newCompletedRows newCompletedCols
        (newRowPlacements, newColPlacements) = newPlacements
    in
        if not updatedFlag
        then (inferredSolution, (newRowPlacements, newColPlacements), newCompletedRows, newCompletedCols)
        else iterativeSolveSeq inferredSolution (newRowPlacements, newColPlacements) rowArgs colArgs newCompletedRows newCompletedCols

iterativeSolvePar :: PartialSolution
               -> PlacementsDict       -- Row and column placements
               -> [Constraint]         -- Row constraints
               -> [Constraint]         -- Column constraints
               -> Set Int              -- Completed localRows
               -> Set Int              -- Completed columns
               -> (PartialSolution, PlacementsDict, Set Int, Set Int)
iterativeSolvePar solutionArray (rowPlacements, colPlacements) rowArgs colArgs completedRows completedCols =
    let
        -- Infer values and update completions
        inferredSolution = inferValuesPar solutionArray (rowPlacements, colPlacements)

        -- Update completions
        (newCompletedRows, newCompletedCols) = updateCompletions inferredSolution

        -- Update placements
        (updatedFlag, newPlacements) = updatePlacements inferredSolution (rowPlacements, colPlacements) newCompletedRows newCompletedCols
        (newRowPlacements, newColPlacements) = newPlacements
    in
        if not updatedFlag
        then (inferredSolution, (newRowPlacements, newColPlacements), newCompletedRows, newCompletedCols)
        else iterativeSolvePar inferredSolution (newRowPlacements, newColPlacements) rowArgs colArgs newCompletedRows newCompletedCols

--------------------------------------------------------------------------------------------------------------------------------------------------------

-- Helper to print a solution in a grid format
printSolution :: PartialSolution -> IO ()
printSolution solution = do
    let ((r1, c1), (r2, c2)) = bounds solution
    mapM_ (putStrLn . concatMap show) [[solution ! (r, c) | c <- [c1..c2]] | r <- [r1..r2]]
    putStrLn "" -- Add a blank line between solutions



-- Generalized solveNonogramFromFile function
solveNonogramFromFile :: (Int -> [Int] -> [[Int]])            -- computeBlocks function
                      -> ([[Int]] -> [Int] -> Int -> [[Int]]) -- generateBlocks function
                      -> (PartialSolution -> PlacementsDict -> [Constraint] -> [Constraint] -> Set Int -> Set Int
                          -> (PartialSolution, PlacementsDict, Set Int, Set Int)) -- iterativeSolve function
                      -> FilePath                            -- File path to Nonogram
                      -> IO ()
solveNonogramFromFile computeBlocksFunc generateBlocksFunc iterativeSolveFunc filePath = do
    -- Parse Nonogram from provided file path
    nonogram <- parseNonogram filePath

    -- Extract row and column arguments
    let rowArgs = rows nonogram
        colArgs = columns nonogram

    putStrLn $ "\nTitle: " ++ title nonogram

{-     putStrLn $ "\nTitle: " ++ title nonogram
    putStrLn "Row Hints (rowArgs):"
    mapM_ print rowArgs
    putStrLn "Column Hints (colArgs):"
    mapM_ print colArgs -}

    -- Compute row and column vector lengths
    let rowVectorLen = length colArgs
    let colVectorLen = length rowArgs

    -- Compute placements using passed-in functions
    let rowPlacements = listArray (0, length rowArgs - 1) $
            map (\arg -> Set.fromList (generateBlocksFunc (computeBlocksFunc rowVectorLen arg) arg rowVectorLen)) rowArgs

    let colPlacements = listArray (0, length colArgs - 1) $
            map (\arg -> Set.fromList (generateBlocksFunc (computeBlocksFunc colVectorLen arg) arg colVectorLen)) colArgs

    -- Initialize partial solution grid
    let partialSolution = array ((0, 0), (length rowArgs - 1, length colArgs - 1))
                        [((r, c), -1) | r <- [0..length rowArgs - 1], c <- [0..length colArgs - 1]]

    let completedRows = Set.empty
    let completedColumns = Set.empty

    -- Step 1: Iterative solving
    -- putStrLn "\nTesting iterativeSolve...\n"
    let placementsDict = (rowPlacements, colPlacements)
    let (finalSolution, _, _, _) =
            iterativeSolveFunc partialSolution placementsDict rowArgs colArgs completedRows completedColumns

    -- putStrLn "\nFinal Solution:"
    printSolution finalSolution

solveNonogramBacktrack :: (Int -> [Int] -> [[Int]])            -- computeBlocks function
                      -> ([[Int]] -> [Int] -> Int -> [[Int]]) -- generateBlocks function
                      -> (PartialSolution -> Array Int (Set [Int]) -> [Constraint] -> [Constraint] -> Set Int -> Set Int
                          -> [PartialSolution] -> [PartialSolution]) -- backtrack function
                      -> FilePath                            -- File path to the Nonogram
                      -> IO ()
solveNonogramBacktrack computeBlocksFunc generateBlocksFunc backtrackFunc filePath = do
    -- Parse Nonogram from provided file path
    nonogram <- parseNonogram filePath

    -- Extract row and column arguments
    let rowArgs = rows nonogram
        colArgs = columns nonogram

    putStrLn $ "\nTitle: " ++ title nonogram

{-     putStrLn $ "\nTitle: " ++ title nonogram
    putStrLn "Row Hints (rowArgs):"
    mapM_ print rowArgs
    putStrLn "Column Hints (colArgs):"
    mapM_ print colArgs -}

    -- Compute row and column vector lengths
    let rowVectorLen = length colArgs
    let colVectorLen = length rowArgs

    -- Compute placements using passed-in functions
    let rowPlacements = listArray (0, length rowArgs - 1) $
            map (\arg -> Set.fromList (generateBlocksFunc (computeBlocksFunc rowVectorLen arg) arg rowVectorLen)) rowArgs

    let colPlacements = listArray (0, length colArgs - 1) $
            map (\arg -> Set.fromList (generateBlocksFunc (computeBlocksFunc colVectorLen arg) arg colVectorLen)) colArgs

    -- Initialize partial solution grid
    let partialSolution = array ((0, 0), (length rowArgs - 1, length colArgs - 1))
                        [((r, c), -1) | r <- [0..length rowArgs - 1], c <- [0..length colArgs - 1]]

    let completedRows = Set.empty
    let completedColumns = Set.empty

    putStrLn "Solving via backtracking..."


    let solutions = backtrack partialSolution rowPlacements rowArgs colArgs completedRows completedColumns []

    putStrLn "Found solutions:"
    mapM_ printSolution solutions



--------------------------------------------------------------------------------------------------------------------------------------------------------


solveSequential :: FilePath -> IO ()
solveSequential = solveNonogramFromFile computeBlocksSeq generateBlocksSeq iterativeSolveSeq

solveParallelComputeBlocks :: FilePath -> IO ()
solveParallelComputeBlocks = solveNonogramFromFile computeBlocksPar generateBlocksSeq iterativeSolveSeq

solveParallelGenerateBlocks :: FilePath -> IO ()
solveParallelGenerateBlocks = solveNonogramFromFile computeBlocksSeq generateBlocksPar iterativeSolveSeq

solveParallelComputeGenerate :: FilePath -> IO ()
solveParallelComputeGenerate = solveNonogramFromFile computeBlocksPar generateBlocksPar iterativeSolveSeq

solveParallelIterativeSolve :: FilePath -> IO ()
solveParallelIterativeSolve = solveNonogramFromFile computeBlocksSeq generateBlocksSeq iterativeSolvePar

solveFullyParallel :: FilePath -> IO ()
solveFullyParallel = solveNonogramFromFile computeBlocksPar generateBlocksPar iterativeSolvePar
