local z,e,K,D,O,n=_G.tonumber,_G.tostring,_G.io,_G.ipairs,_G.string,_G.table local t,b=setmetatable,getmetatable local P=require("\112\108\046\116\097\098\108\101\120")local A=require("\112\108\046\117\116\105\108\115")local c=require("\112\108\046\116\121\112\101\115")local f,a,U,T,p,i,B=P.imap,P.map,P.reduce,P.keys,P.map2,P.set,P.index_by local X=n.remove local F,Y,V=A.splitv,A.fprintf,A.assert_arg local k=O.byte local l=K.stdout local y=math.min local N={}local function R(z,e)local K=b(z)if K then t(e,K)end return e end local function Q(z)return t(z,require("\112\108\046\076\105\115\116"))end function N.size(z)V(1,z,"\116\097\098\108\101")return#z,#z[1]end do local function z(z,e)return z[e]end function N.column(e,K)V(1,e,"\116\097\098\108\101")return Q(f(z,e,K))end end local h=N.column function N.row(z,e)V(1,z,"\116\097\098\108\101")local K=z[e]local O={}for z,e in D(K)do O[z]=e end return Q(O)end function N.map(z,e,K)V(2,e,"\116\097\098\108\101")z=A.function_arg(1,z)return R(e,f(function(e)return f(z,e,K)end,e))end function N.reduce_rows(z,e)V(1,e,"\116\097\098\108\101")return a(function(e)return U(z,e)end,e)end function N.reduce_cols(z,e)V(1,e,"\116\097\098\108\101")return a(function(K)return U(z,h(e,K))end,T(e[1]))end function N.reduce2(z,e,K)V(3,K,"\116\097\098\108\101")local D=N.reduce_rows(e,K)return U(z,D)end function N.map2(z,e,K,D,O,n)V(1,D,"\116\097\098\108\101")V(2,O,"\116\097\098\108\101")z=A.function_arg(1,z)if e==1 and K==2 then return f(function(e)return p(z,D,e,n)end,O)elseif e==2 and K==1 then return f(function(e)return p(z,e,O,n)end,D)elseif e==1 and K==1 then return p(z,D,O)elseif e==2 and K==2 then return p(function(e,K)return p(z,e,K,n)end,D,O)end end function N.product(z,e,K)z=A.function_arg(1,z)V(2,e,"\116\097\098\108\101")V(3,K,"\116\097\098\108\101")local O={}for K,D in D(K)do O[K]=a(z,e,D)end return O end function N.flatten(z)local e={}local K=1 local D,O=N.size(z)for D=1,D,1 do local n=z[D]for z=1,O,1 do e[K]=n[z]K=K+1 end end return Q(e)end function N.reshape(z,e,K)local D,O=N.size(z)local n=(D*O)/e local t={}local b,P=1,1 for e=1,e,1 do local A={}for e=1,n,1 do A[e]=z[b][P]if not K then P=P+1 if P>O then b=b+1 P=1 end else b=b+1 if b>D then P=P+1 b=1 end end end t[e]=A end return R(z,t)end function N.transpose(z)V(1,z,"\116\097\098\108\101")local e,K=N.size(z)return N.reshape(z,K,true)end function N.swap_rows(z,e,K)V(1,z,"\116\097\098\108\101")z[e],z[K]=z[K],z[e]return z end function N.swap_cols(z,e,K)V(1,z,"\116\097\098\108\101")for z,D in D(z)do D[e],D[K]=D[K],D[e]end return z end function N.extract_rows(z,e)return R(z,B(z,e))end function N.extract_cols(z,e)V(1,z,"\116\097\098\108\101")local K={}for D=1,#z,1 do K[D]=B(z[D],e)end return R(z,K)end N.remove_row=X function N.remove_col(z,e)V(1,z,"\116\097\098\108\101")for K=1,#z,1 do X(z[K],e)end end do local function e(e)local K,D=e:match("\082\040\037\100\043\041\067\040\037\100\043\041")if K then K,D=z(K),z(D)return K,D end D,K=e:match("\040\037\097\043\041\040\037\100\043\041")if D then local e=0 for z=1,#D,1 do e=((e*26+k(D:sub(z,z)))-k("\065"))+1 end return z(K),e end error("\098\097\100 \099\101\108\108 \115\112\101\099\105\102\105\101\114\058 "..e)end function N.parse_range(z)V(1,z,"\115\116\114\105\110\103")z=z:upper()if z:find("\058")then local K,D=F(z,"\058")local O,n=e(K)local t,b=e(D)return O,n,t,b else local K,D=e(z)return K,D end end end function N.range(...)return N.slice(...)end local q do local function z(z,e)if not z then return z end if z<0 then z=(e+z)+1 end if z<1 then z=1 end if z>e then z=e end return z end function N.default_range(e,K,D,O,n)if type(K)=="\115\116\114\105\110\103"and not(D or O or n)then K,D,O,n=N.parse_range(K)end local t,b=N.size(e)K=z(K or 1,t)D=z(D or 1,b)O=z(O or t,t)n=z(n or b,b)return K,D,O,n end q=N.default_range end function N.slice(z,e,K,D,O)V(1,z,"\116\097\098\108\101")e,K,D,O=q(z,e,K,D,O)local n={}for e=e,D,1 do local D local t=z[e]if K==O then D=t[K]else D={}for z=K,O,1 do D[#D+1]=t[z]end end n[#n+1]=D end if e==D then n=n[1]end return R(z,n)end function N.set(z,e,K,D,O,n)K,D,O,n=q(z,K,D,O,n)local t=K if c.is_callable(e)then local z=e e=function(e)return z(t,e)end end while t<=O do i(z[t],e,D,n)t=t+1 end end function N.write(z,K,D,O,n,t,b)V(1,z,"\116\097\098\108\101")K=K or l local P if D then P=function(z,e)Y(K,D,z[e])end else P=function(z,D)K:write(e(z[D])," ")end end local function A()K:write("\n")end N.forall(z,P,A,O,n,t,b)end function N.forall(z,e,K,D,O,n,t)V(1,z,"\116\097\098\108\101")D,O,n,t=q(z,D,O,n,t)for D=D,n,1 do local n=z[D]for z=O,t,1 do e(n,z)end if K then K(D)end end end function N.move(z,e,K,D,O,n,t,b)V(1,z,"\116\097\098\108\101")V(4,D,"\116\097\098\108\101")O,n,t,b=q(D,O,n,t,b)local P,A=N.size(z)t,b=y(P,t),y(A,b)K=K-1 for O=O,t,1 do local t,P=z[(O+e)-1],D[O]for z=n,b,1 do t[z+K]=P[z]end end end function N.iter(z,e,K,D,O,n)V(1,z,"\116\097\098\108\101")K,D,O,n=q(z,K,D,O,n)local t,b=K,D-1 local P=z[t]return function()b=b+1 if b>n then b=D t=t+1 P=z[t]if t>O then return nil end end if e then return t,b,P[b]else return P[b]end end end function N.columns(z)V(1,z,"\116\097\098\108\101")local e=#z[1]local K=0 return function()K=K+1 if K>e then return nil end return h(z,K),K end end function N.rows(z)V(1,z,"\116\097\098\108\101")local e=#z local K=0 return function()K=K+1 if K>e then return nil end return N.row(z,K),K end end function N.new(z,e,K)local D={}local O=c.is_callable(K)for z=1,z,1 do local n={}if O then for e=1,e,1 do n[e]=K(z,e)end else for z=1,e,1 do n[z]=K end end D[z]=n end return D end return N
--- Operations on two-dimensional arrays.
-- See @{02-arrays.md.Operations_on_two_dimensional_tables|The Guide}
--
-- The size of the arrays is determined by using the length operator `#` hence
-- the module is not `nil` safe, and the usual precautions apply.
--
-- Note: all functions taking `i1,j1,i2,j2` as arguments will normalize the
-- arguments using `default_range`.
--
-- Dependencies: `pl.utils`, `pl.tablex`, `pl.types`
-- @module pl.array2d

local tonumber,tostring,io,ipairs,string,table =
    _G.tonumber,_G.tostring,_G.io,_G.ipairs,_G.string,_G.table
local setmetatable,getmetatable = setmetatable,getmetatable

local tablex = require 'pl.tablex'
local utils = require 'pl.utils'
local types = require 'pl.types'
local imap,tmap,reduce,keys,tmap2,tset,index_by = tablex.imap,tablex.map,tablex.reduce,tablex.keys,tablex.map2,tablex.set,tablex.index_by
local remove = table.remove
local splitv,fprintf,assert_arg = utils.splitv,utils.fprintf,utils.assert_arg
local byte = string.byte
local stdout = io.stdout
local min = math.min


local array2d = {}

local function obj (int,out)
    local mt = getmetatable(int)
    if mt then
        setmetatable(out,mt)
    end
    return out
end

local function makelist (res)
    return setmetatable(res, require('pl.List'))
end

--- return the row and column size.
-- Size is calculated using the Lua length operator #, so usual precautions
-- regarding `nil` values apply.
-- @array2d a a 2d array
-- @treturn int number of rows (`#a`)
-- @treturn int number of cols (`#a[1]`)
function array2d.size (a)
    assert_arg(1,a,'table')
    return #a,#a[1]
end

do
    local function index (t,k)
        return t[k]
    end

    --- extract a column from the 2D array.
    -- @array2d a 2d array
    -- @param j column index
    -- @return 1d array
    function array2d.column (a,j)
        assert_arg(1,a,'table')
        return makelist(imap(index,a,j))
    end
end
local column = array2d.column

--- extract a row from the 2D array.
-- Added in line with `column`, for read-only purposes directly
-- accessing a[i] is more performant.
-- @array2d a 2d array
-- @param i row index
-- @return 1d array (copy of the row)
function array2d.row(a,i)
    assert_arg(1,a,'table')
    local row = a[i]
    local r = {}
    for n,v in ipairs(row) do
        r[n] = v
    end
    return makelist(r)
end

--- map a function over a 2D array
-- @func f a function of at least one argument
-- @array2d a 2d array
-- @param arg an optional extra argument to be passed to the function.
-- @return 2d array
function array2d.map (f,a,arg)
    assert_arg(2,a,'table')
    f = utils.function_arg(1,f)
    return obj(a,imap(function(row) return imap(f,row,arg) end, a))
end

--- reduce the rows using a function.
-- @func f a binary function
-- @array2d a 2d array
-- @return 1d array
-- @see pl.tablex.reduce
function array2d.reduce_rows (f,a)
    assert_arg(1,a,'table')
    return tmap(function(row) return reduce(f,row) end, a)
end

--- reduce the columns using a function.
-- @func f a binary function
-- @array2d a 2d array
-- @return 1d array
-- @see pl.tablex.reduce
function array2d.reduce_cols (f,a)
    assert_arg(1,a,'table')
    return tmap(function(c) return reduce(f,column(a,c)) end, keys(a[1]))
end

--- reduce a 2D array into a scalar, using two operations.
-- @func opc operation to reduce the final result
-- @func opr operation to reduce the rows
-- @param a 2D array
function array2d.reduce2 (opc,opr,a)
    assert_arg(3,a,'table')
    local tmp = array2d.reduce_rows(opr,a)
    return reduce(opc,tmp)
end

--- map a function over two arrays.
-- They can be both or either 2D arrays
-- @func f function of at least two arguments
-- @int ad order of first array (`1` if `a` is a list/array, `2` if it is a 2d array)
-- @int bd order of second array (`1` if `b` is a list/array, `2` if it is a 2d array)
-- @tab a 1d or 2d array
-- @tab b 1d or 2d array
-- @param arg optional extra argument to pass to function
-- @return 2D array, unless both arrays are 1D
function array2d.map2 (f,ad,bd,a,b,arg)
    assert_arg(1,a,'table')
    assert_arg(2,b,'table')
    f = utils.function_arg(1,f)
    if ad == 1 and bd == 2 then
        return imap(function(row)
            return tmap2(f,a,row,arg)
        end, b)
    elseif ad == 2 and bd == 1 then
        return imap(function(row)
            return tmap2(f,row,b,arg)
        end, a)
    elseif ad == 1 and bd == 1 then
        return tmap2(f,a,b)
    elseif ad == 2 and bd == 2 then
        return tmap2(function(rowa,rowb)
            return tmap2(f,rowa,rowb,arg)
        end, a,b)
    end
end

--- cartesian product of two 1d arrays.
-- @func f a function of 2 arguments
-- @array t1 a 1d table
-- @array t2 a 1d table
-- @return 2d table
-- @usage product('..',{1,2},{'a','b'}) == {{'1a','2a'},{'1b','2b'}}
function array2d.product (f,t1,t2)
    f = utils.function_arg(1,f)
    assert_arg(2,t1,'table')
    assert_arg(3,t2,'table')
    local res = {}
    for i,v in ipairs(t2) do
        res[i] = tmap(f,t1,v)
    end
    return res
end

--- flatten a 2D array.
-- (this goes over columns first.)
-- @array2d t 2d table
-- @return a 1d table
-- @usage flatten {{1,2},{3,4},{5,6}} == {1,2,3,4,5,6}
function array2d.flatten (t)
    local res = {}
    local k = 1
    local rows, cols = array2d.size(t)
    for r = 1, rows do
        local row = t[r]
        for c = 1, cols do
            res[k] = row[c]
            k = k + 1
        end
    end
    return makelist(res)
end

--- reshape a 2D array. Reshape the aray by specifying a new nr of rows.
-- @array2d t 2d array
-- @int nrows new number of rows
-- @bool co use column-order (Fortran-style) (default false)
-- @return a new 2d array
function array2d.reshape (t,nrows,co)
    local nr,nc = array2d.size(t)
    local ncols = nr*nc / nrows
    local res = {}
    local ir,ic = 1,1
    for i = 1,nrows do
        local row = {}
        for j = 1,ncols do
            row[j] = t[ir][ic]
            if not co then
                ic = ic + 1
                if ic > nc then
                    ir = ir + 1
                    ic = 1
                end
            else
                ir = ir + 1
                if ir > nr then
                    ic = ic + 1
                    ir = 1
                end
            end
        end
        res[i] = row
    end
    return obj(t,res)
end

--- transpose a 2D array.
-- @array2d t 2d array
-- @return a new 2d array
function array2d.transpose(t)
  assert_arg(1,t,'table')
  local _, c = array2d.size(t)
  return array2d.reshape(t,c,true)
end

--- swap two rows of an array.
-- @array2d t a 2d array
-- @int i1 a row index
-- @int i2 a row index
-- @return t (same, modified 2d array)
function array2d.swap_rows (t,i1,i2)
    assert_arg(1,t,'table')
    t[i1],t[i2] = t[i2],t[i1]
    return t
end

--- swap two columns of an array.
-- @array2d t a 2d array
-- @int j1 a column index
-- @int j2 a column index
-- @return t (same, modified 2d array)
function array2d.swap_cols (t,j1,j2)
    assert_arg(1,t,'table')
    for _, row in ipairs(t) do
        row[j1],row[j2] = row[j2],row[j1]
    end
    return t
end

--- extract the specified rows.
-- @array2d t 2d array
-- @tparam {int} ridx a table of row indices
-- @return a new 2d array with the extracted rows
function array2d.extract_rows (t,ridx)
    return obj(t,index_by(t,ridx))
end

--- extract the specified columns.
-- @array2d t 2d array
-- @tparam {int} cidx a table of column indices
-- @return a new 2d array with the extracted colums
function array2d.extract_cols (t,cidx)
    assert_arg(1,t,'table')
    local res = {}
    for i = 1,#t do
        res[i] = index_by(t[i],cidx)
    end
    return obj(t,res)
end

--- remove a row from an array.
-- @function array2d.remove_row
-- @array2d t a 2d array
-- @int i a row index
array2d.remove_row = remove

--- remove a column from an array.
-- @array2d t a 2d array
-- @int j a column index
function array2d.remove_col (t,j)
    assert_arg(1,t,'table')
    for i = 1,#t do
        remove(t[i],j)
    end
end

do
    local function _parse (s)
        local r, c = s:match 'R(%d+)C(%d+)'
        if r then
            r,c = tonumber(r),tonumber(c)
            return r,c
        end
        c,r = s:match '(%a+)(%d+)'
        if c then
            local cv = 0
            for i = 1, #c do
              cv = cv * 26 + byte(c:sub(i,i)) - byte 'A' + 1
            end
            return tonumber(r), cv
        end
        error('bad cell specifier: '..s)
    end

    --- parse a spreadsheet range or cell.
    -- The range/cell can be specified either as 'A1:B2' or 'R1C1:R2C2' or for
    -- single cells as 'A1' or 'R1C1'.
    -- @string s a range (case insensitive).
    -- @treturn int start row
    -- @treturn int start col
    -- @treturn int end row (or `nil` if the range was a single cell)
    -- @treturn int end col (or `nil` if the range was a single cell)
    function array2d.parse_range (s)
        assert_arg(1,s,'string')
        s = s:upper()
        if s:find ':' then
            local start,finish = splitv(s,':')
            local i1,j1 = _parse(start)
            local i2,j2 = _parse(finish)
            return i1,j1,i2,j2
        else -- single value
            local i,j = _parse(s)
            return i,j
        end
    end
end

--- get a slice of a 2D array.
-- Same as `slice`.
-- @see slice
function array2d.range (...)
    return array2d.slice(...)
end

local default_range do
    local function norm_value(v, max)
        if not v then return v end
        if v < 0 then
            v = max + v + 1
        end
        if v < 1 then v = 1 end
        if v > max then v = max end
        return v
    end

    --- normalizes coordinates to valid positive entries and defaults.
    -- Negative indices will be counted from the end, too low, or too high
    -- will be limited by the array sizes.
    -- @array2d t a 2D array
    -- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
    -- @tparam[opt=1] int j1 start col
    -- @tparam[opt=N] int i2 end row
    -- @tparam[opt=M] int j2 end col
    -- @see parse_range
    -- @return i1, j1, i2, j2
    function array2d.default_range (t,i1,j1,i2,j2)
        if (type(i1) == 'string') and not (j1 or i2 or j2) then
            i1, j1, i2, j2 = array2d.parse_range(i1)
        end
        local nr, nc = array2d.size(t)
        i1 = norm_value(i1 or 1, nr)
        j1 = norm_value(j1 or 1, nc)
        i2 = norm_value(i2 or nr, nr)
        j2 = norm_value(j2 or nc, nc)
        return i1,j1,i2,j2
    end
    default_range = array2d.default_range
end

--- get a slice of a 2D array. Note that if the specified range has
-- a 1D result, the rank of the result will be 1.
-- @array2d t a 2D array
-- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
-- @tparam[opt=1] int j1 start col
-- @tparam[opt=N] int i2 end row
-- @tparam[opt=M] int j2 end col
-- @see parse_range
-- @return an array, 2D in general but 1D in special cases.
function array2d.slice (t,i1,j1,i2,j2)
    assert_arg(1,t,'table')
    i1,j1,i2,j2 = default_range(t,i1,j1,i2,j2)
    local res = {}
    for i = i1,i2 do
        local val
        local row = t[i]
        if j1 == j2 then
            val = row[j1]
        else
            val = {}
            for j = j1,j2 do
                val[#val+1] = row[j]
            end
        end
        res[#res+1] = val
    end
    if i1 == i2 then res = res[1] end
    return obj(t,res)
end

--- set a specified range of an array to a value.
-- @array2d t a 2D array
-- @param value the value (may be a function, called as `val(i,j)`)
-- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
-- @tparam[opt=1] int j1 start col
-- @tparam[opt=N] int i2 end row
-- @tparam[opt=M] int j2 end col
-- @see parse_range
-- @see tablex.set
function array2d.set (t,value,i1,j1,i2,j2)
    i1,j1,i2,j2 = default_range(t,i1,j1,i2,j2)
    local i = i1
    if types.is_callable(value) then
        local old_f = value
        value = function(j)
            return old_f(i,j)
        end
    end
    while i <= i2 do
        tset(t[i],value,j1,j2)
        i = i + 1
    end
end

--- write a 2D array to a file.
-- @array2d t a 2D array
-- @param f a file object (default stdout)
-- @string fmt a format string (default is just to use tostring)
-- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
-- @tparam[opt=1] int j1 start col
-- @tparam[opt=N] int i2 end row
-- @tparam[opt=M] int j2 end col
-- @see parse_range
function array2d.write (t,f,fmt,i1,j1,i2,j2)
    assert_arg(1,t,'table')
    f = f or stdout
    local rowop
    if fmt then
        rowop = function(row,j) fprintf(f,fmt,row[j]) end
    else
        rowop = function(row,j) f:write(tostring(row[j]),' ') end
    end
    local function newline()
        f:write '\n'
    end
    array2d.forall(t,rowop,newline,i1,j1,i2,j2)
end

--- perform an operation for all values in a 2D array.
-- @array2d t 2D array
-- @func row_op function to call on each value; `row_op(row,j)`
-- @func end_row_op function to call at end of each row; `end_row_op(i)`
-- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
-- @tparam[opt=1] int j1 start col
-- @tparam[opt=N] int i2 end row
-- @tparam[opt=M] int j2 end col
-- @see parse_range
function array2d.forall (t,row_op,end_row_op,i1,j1,i2,j2)
    assert_arg(1,t,'table')
    i1,j1,i2,j2 = default_range(t,i1,j1,i2,j2)
    for i = i1,i2 do
        local row = t[i]
        for j = j1,j2 do
            row_op(row,j)
        end
        if end_row_op then end_row_op(i) end
    end
end

---- move a block from the destination to the source.
-- @array2d dest a 2D array
-- @int di start row in dest
-- @int dj start col in dest
-- @array2d src a 2D array
-- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
-- @tparam[opt=1] int j1 start col
-- @tparam[opt=N] int i2 end row
-- @tparam[opt=M] int j2 end col
-- @see parse_range
function array2d.move (dest,di,dj,src,i1,j1,i2,j2)
    assert_arg(1,dest,'table')
    assert_arg(4,src,'table')
    i1,j1,i2,j2 = default_range(src,i1,j1,i2,j2)
    local nr,nc = array2d.size(dest)
    i2, j2 = min(nr,i2), min(nc,j2)
    --i1, j1 = max(1,i1), max(1,j1)
    dj = dj - 1
    for i = i1,i2 do
        local drow, srow = dest[i+di-1], src[i]
        for j = j1,j2 do
            drow[j+dj] = srow[j]
        end
    end
end

--- iterate over all elements in a 2D array, with optional indices.
-- @array2d a 2D array
-- @bool indices with indices (default false)
-- @tparam[opt=1] int|string i1 start row or spreadsheet range passed to `parse_range`
-- @tparam[opt=1] int j1 start col
-- @tparam[opt=N] int i2 end row
-- @tparam[opt=M] int j2 end col
-- @see parse_range
-- @return either `value` or `i,j,value` depending on the value of `indices`
function array2d.iter(a,indices,i1,j1,i2,j2)
    assert_arg(1,a,'table')
    i1,j1,i2,j2 = default_range(a,i1,j1,i2,j2)
    local i,j = i1,j1-1
    local row = a[i]
    return function()
        j = j + 1
        if j > j2 then
            j = j1
            i = i + 1
            row = a[i]
            if i > i2 then
                return nil
            end
        end
        if indices then
            return i,j,row[j]
        else
            return row[j]
        end
    end
end

--- iterate over all columns.
-- @array2d a a 2D array
-- @return column, column-index
function array2d.columns(a)
  assert_arg(1,a,'table')
  local n = #a[1]
  local i = 0
  return function()
      i = i + 1
      if i > n then return nil end
      return column(a,i), i
  end
end

--- iterate over all rows.
-- Returns a copy of the row, for read-only purposes directly iterating
-- is more performant; `ipairs(a)`
-- @array2d a a 2D array
-- @return row, row-index
function array2d.rows(a)
  assert_arg(1,a,'table')
  local n = #a
  local i = 0
  return function()
      i = i + 1
      if i > n then return nil end
      return array2d.row(a,i), i
  end
end

--- new array of specified dimensions
-- @int rows number of rows
-- @int cols number of cols
-- @param val initial value; if it's a function then use `val(i,j)`
-- @return new 2d array
function array2d.new(rows,cols,val)
    local res = {}
    local fun = types.is_callable(val)
    for i = 1,rows do
        local row = {}
        if fun then
            for j = 1,cols do row[j] = val(i,j) end
        else
            for j = 1,cols do row[j] = val end
        end
        res[i] = row
    end
    return res
end

return array2d
