module SymmetrySnap
    using LinearAlgebra
    using Random
    using DelimitedFiles
    using MinkowskiReduction

    """
        getLatticeDiff(oldA, a, oldB, b, oldC, c, output)

    Save the change in lattice to output vector

    # Arguments
    - `oldA, oldB, oldC::Vector{Float64}`: Old lattice vectors
    - `a, b, c::Vector{Float64}`: New lattice vectors
    - `output::Vector{String}`: Output list
    """
    function getLatticeDiff(oldA, a, oldB, b, oldC, c, numSymmetries, output)
        push!(output, "Old Lattice -> New Lattice : L2 Norm, SymOps = $(numSymmetries)")
        push!(output, join([oldA, a, norm(oldA - a)]," -> "," : "))
        push!(output, join([oldB, b, norm(oldB - b)]," -> "," : "))
        push!(output, join([oldC, c, norm(oldC - c)]," -> "," : "))
        push!(output, "")
        return output
    end

    """
        getAtomDiff(oldAtomPos, atomPos, symmetries, output)

    Save the change in atom positions and number of symmetries found to output vector

    # Arguments
    - `oldAtomPos::Vector{Vector{Float64}}`: Old atom positions
    - `atomPos::Vector{Vector{Float64}}`: New atom positions
    - `symmetries::Int8`: Number of symmetries found
    - `output::Vector{String}`: Output list
    """
    function getAtomDiff(oldAtomPos, atomPos, symmetries, output)
        push!(output, "Old Atom Position -> New Atom Position : L2 Norm, SymOps")
        for i in 1:length(atomPos)
            push!(output, join([join([oldAtomPos[i], atomPos[i], norm(oldAtomPos[i] - atomPos[i])]," -> "," : "),symmetries[i]],", "))
        end
        return output
    end

    """
        mapAtomsToCell(a, b, c, latVecInv, atomPositions)

    Map given atom positions into positions inside the unit cell.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: Inverse of lattice vectors
    - `atomPositions::Vector{Vector{Float64}}`

    # Examples
    ```jldoctest
    julia> mapAtomsToCell([1,0,0], [0,1,0], [0,0,1], [1, 0, 0; 0, 1, 0; 0, 0, 1], [1.5,-0.5,3.5])
    [0.5, 0.5, 0.5]
    ```
    """
    function mapAtomsToCell(a, b, c, latVecInv, atomPositions)
        newPositions = []
        for atom in atomPositions
            addTo = true
            shift = latVecInv * atom
            for i in 1:3
                shift[i] = floor(shift[i])
            end
            newPos = atom - a*shift[1] - b*shift[2] - c*shift[3]
            for newAtom in newPositions
                if isapprox(newPos, newAtom)
                    addTo = false
                    break
                end
            end
            if addTo == true
                push!(newPositions, newPos)
            end
        end
        return newPositions
    end

    """
        toDirect(a, b, c, atomPositions)

    Convert atom positions in Cartesian coordinates to direct coordinates.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `atomPositions::Vector{Vector{Float64}}`: atom positions in Cartesian coordinates

    # Examples


    """
    function toDirect(a, b, c, atomPositions)
        latVecs = Array{Float64, 2}(undef, 3, 3)
        latVecs[1:end,1] = a
        latVecs[1:end,2] = b
        latVecs[1:end,3] = c

        newPositions = Vector{Vector{Float64}}(undef, length(atomPositions))
        for i in 1:length(atomPositions)
            newPositions[i] = \(latVecs, atomPositions[i])
        end
        return newPositions
    end

    """
        toCartesian(a, b, c, atomPositions)

        Convert atom positions in direct coordinates to Cartesian coordinates.

        # Arguments
        - `a,b,c::Vector{Float64}`: lattice vectors
        - `atomPositions::Vector{Vector{Float64}}`: atom positions in direct coordinates

        # Examples


    """
    function toCartesian(a, b, c, atomPositions)
        latVecs = Array{Float64, 2}(undef, 3, 3)
        latVecs[1:end,1] = a
        latVecs[1:end,2] = b
        latVecs[1:end,3] = c

        newPositions = Vector{Vector{Float64}}(undef, length(atomPositions))
        for i in 1:length(atomPositions)
            newPositions[i] = latVecs * atomPositions[i]
        end
        return newPositions
    end

    """
        test_atom(a, b, c, lat_vec_inv, symmetry_op, atom_1_pos, atom_1_index, atom_2_pos, atom_2_index, natoms[, epsilon=1e-2])

    Test if atom 1 maps to the position of atom 2 when the symmetry operator is applied. If valid, returns matrix A that represents the mapping from atom 1 to atom 2.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix
    - `symmetryOp::Array{Float64, 2}`: symmetry operator
    - `atom_1_index, atom_2_index::{Int}`: indexes of atoms being compared
    - `atom_1_pos, atom_2_pos::Vector{Float64}`: positions of atoms being compared
    - `nAtoms::Int8`: Total number of atoms in the system
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function testAtom(a, b, c, latVecInv, symmetryOp, atom1Pos, atom1Index, atom2Pos, atom2Index, nAtoms, shift, epsilon=1e-2)
        # In the case of a different crystal structure, A would be the lattice vectors as rows
        m = symmetryOp * atom1Pos #what is m??

        if shift != nothing
            m = m .+ shift
        end

        compareCell = latVecInv * m
        for i in 1:3
            compareCell[i] = floor(compareCell[i])
        end

        transformedAtom = symmetryOp * atom1Pos .- a*compareCell[1] .- b*compareCell[2] .- c*compareCell[3]
        if shift == nothing
            shift = atom2Pos .- transformedAtom
        end

        transformedAtom = transformedAtom .+ shift

        if norm(atom2Pos .- symmetryOp * atom1Pos .- shift) < epsilon || norm(atom2Pos .- transformedAtom) < epsilon
            A = zeros(Float64, 3, 3*nAtoms+1)
            if norm(atom2Pos .- symmetryOp * atom1Pos .- shift) < epsilon
                compareCell = [0,0,0]
            end
            for i in 1:3
                for j in 1:3
                    if symmetryOp[i,j] != 0
                        A[i,3*(atom1Index-1)+j] += symmetryOp[i,j]
                    end
                    if -Matrix(I, 3, 3)[i,j] != 0
                        A[i,3*(atom2Index-1)+j] += -Matrix(I, 3, 3)[i,j]
                    end
                end
            end
            A[:,3*nAtoms+1] = a*compareCell[1] .+ b*compareCell[2] .+ c*compareCell[3] .- shift
            A[abs.(A) .< 1e-6] .= 0
            if isapprox(A, zero(A), atol = 1e-10)
                return nothing, shift, true
            end
            return A, shift, true
        else
            return nothing, shift, false
        end
    end

    """
        runSymmetry(a,b,c,latVecInv, symmetryOp, atomPositions[, epsilon=1e-2])

    Test the given symmetry operator on each atom in atomPositions to verify if it is a true symmetry, and return the list of symmetry equivalent linear algebra problems.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix
    - `symmetryOp::Array{Float64, 2}`: symmetry operator
    - `atomPositions::Vector{Vector{Float64}}`: atom positions
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function runSymmetry(a, b, c, latVecInv, symmetryOp, atomPositions, epsilon=1e-2)
        # Which atom to shift by
        for i in 1:length(atomPositions)
            As = Vector{Array{Float64, 2}}(undef, 0)
            global shift = nothing
            indexFound = []
            # Iterate over all atoms for atom 1
            for j in 1:length(atomPositions)
                global symmFound = false
                # Iterate over all atoms for atom 2
                for k in 1:length(atomPositions)
                    if i+k-1 in indexFound
                        continue
                    end
                    A, shift, symmFound = testAtom(a, b, c, latVecInv, symmetryOp, atomPositions[j], j, atomPositions[(i+k-2)%length(atomPositions)+1], (i+k-2)%length(atomPositions)+1, length(atomPositions), shift, epsilon)
                    if A != nothing
                        push!(As, A)
                        push!(indexFound, i+k-1)
                    end
                    if symmFound
                        push!(indexFound, i+k-1)
                        break
                    end
                end
                if !symmFound
                    break
                end
            end
            if symmFound
                return As
            end
        end
        return nothing
    end

    """
        symmetrize(a, b, c, lat_vec_inv, atom_positions[, epsilon=1e-2])

    Using the valid symmetry operators found, solve the \\(Ax=b\\) problem to get the smallest correction in atom positions needed to align with the improved lattice vectors.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix
    - `atomPositions::Vector{Vector{Float64}}`: atom positions
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function symmetrize(a, b, c, latVecInv, atomPositions, symmetries, epsilon=1e-2)
        As = Vector{Array{Float64, 2}}(undef, 0)
        symm_num = 0
        for i in 1:length(symmetries)
            newA = runSymmetry(a, b, c, latVecInv, symmetries[i], atomPositions, epsilon)
            if newA != nothing
                symm_num += 1
                check = false
                for new in newA
                    for A in As
                        if isapprox(A,new, atol = 1e-10)
                            check = true
                            break
                        end
                    end
                    if !check
                        push!(As,new)
                    end
                end
            end
        end
        if length(As) == 0
            return atomPositions, zeros(length(atomPositions))
        end

        A = zeros(Float64, 3*length(As), 3*length(atomPositions))
        B = zeros(Float64, 3*length(As))
        for i in 0:length(As)-1
            A[3*i+1:3*i+3,:] = As[i+1][1:end,1:end .!= length(As[i+1][1,1:end])]
            B[3*i+1:3*i+3] = As[i+1][1:end,end]
        end

        aSVD = svd(A)
        newA = zeros(3*length(As), 3*length(atomPositions))
        for i in 1:length(aSVD.S)
            uArray = aSVD.U[1:end, i]
            vArray = aSVD.Vt[i, 1:end]
            part = aSVD.S[i]*(aSVD.U[1:end, i]*transpose(aSVD.Vt[i, 1:end]))
            part[abs.(part) .< 1e-6] .= 0
            newA += part
        end
        A = newA

        xUser = vcat(atomPositions...)
        xDiff = xUser - \(A,B)
        null = transpose(LinearAlgebra.nullspace(A))

        for i in 1:length(null[1:end, 1])
            xDiff -= null[i,1:end]*dot(xDiff,null[i,1:end])/norm(null[i,1:end])
        end
        xCorrected = xUser - xDiff
        newXs = Array{Array{Float64, 1}, 1}(undef, length(atomPositions))
        for i in 1:length(atomPositions)
            newX = Array{Float64, 1}(undef, 3)
            for j in 1:3
                newX[j] = xCorrected[3*(i-1) + j]
            end
            newXs[i] = newX
        end

        return a, b, c, newXs, repeat([symm_num], length(atomPositions))
    end

    """
        getSymmetryOps(a,b,c,latVecInv)

    Finds the applicable symmetry operators for the system by solving the following equation,

    \\(R = A B^{-1}\\),

    where \\(A\\) is the original lattice vectors, \\(B\\) is the new lattice vectors, and \\(R\\) is the symmetry operator being tested.

    If \\(R R^T = I\\), \\(R\\) is a valid symmetry operator.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix

    # Examples

    """
    function getSymmetryOps(a, b, c, latVecInv, atomPositions, epsilon=1e-2)
        translations = [[0,0,1],[0,1,0],[1,0,0],[0,0,-1],[0,-1,0],[-1,0,0],[1,1,0],[1,-1,0],[-1,1,0],
                        [-1,-1,0],[1,0,1],[1,0,-1],[-1,0,1],[-1,0,-1],[0,1,1],[0,1,-1],[0,-1,1],[0,-1,-1],
                        [1,1,1],[1,1,-1],[1,-1,1],[-1,1,1],[1,-1,-1],[-1,1,-1],[-1,-1,1],[-1,-1,-1]]
        latticePoints = []
        for v in translations
            push!(latticePoints,v[1]*a+v[2]*b+v[3]*c)
        end

        oldLattice = Array{Float64, 2}(undef, 3, 3)
        oldLattice[1:end, 1] = a
        oldLattice[1:end, 2] = b
        oldLattice[1:end, 3] = c

        newLattices = Vector{Array{Float64, 2}}(undef, 15600)
        transformations = Vector{Array{Float64, 2}}(undef, 15600)
        index = 1
        for i in 1:length(latticePoints)
            for j in 1:length(latticePoints)
                for k in 1:length(latticePoints)
                    if i == k || j == k || i == j
                        continue
                    end
                    newLattice = Array{Float64, 2}(undef, 3, 3)
                    newLattice[1:end, 1] = latticePoints[i]
                    newLattice[1:end, 2] = latticePoints[j]
                    newLattice[1:end, 3] = latticePoints[k]
                    newLattices[index] = newLattice

                    transformLattice = Array{Float64, 2}(undef, 3, 3)
                    transformLattice[1, 1:end] = translations[i]
                    transformLattice[2, 1:end] = translations[j]
                    transformLattice[3, 1:end] = translations[k]
                    transformations[index] = transformLattice
                    index += 1
                end
            end
        end

        allOperators = []
        allErrors = []
        for i in 1:length(newLattices)
            operator = newLattices[i]*latVecInv
            operatorT = transpose(operator)
            push!(allOperators, operator)
            push!(allErrors, norm(operator*operatorT - Matrix(I,3,3)))
        end

        p = sortperm(allErrors)
        sortOperators = allOperators[p]
        sortErrors = allErrors[p]
        sortTransforms = transformations[p]

        # Symmetrix Matrix, so only 6 elements: xx, yy, zz, xy, xz, yz
        elementMap = Dict(
            (1, 1) => 1,
            (2, 2) => 2,
            (3, 3) => 3,
            (1, 2) => 4,
            (2, 1) => 4,
            (1, 3) => 5,
            (3, 1) => 5,
            (2, 3) => 6,
            (3, 2) => 6
            )
        numOps = [2, 3, 4, 6, 8, 12, 16, 24, 48]
        lattices = []
        errors = []
        opCounts = []
        for numOp in numOps
            A = zeros(6*numOp, 6)
            for opNum in 0:numOp-1
                for i in 1:3
                    for j in i:3
                        A[6*opNum + elementMap[(i, j)], elementMap[(i, j)]] -= 1
                        for k in 1:3
                            for l in 1:3
                                A[6*opNum + elementMap[(i, j)], elementMap[(k, l)]] += sortTransforms[opNum+1][i, k]*sortTransforms[opNum+1][j, l]
                            end
                        end
                    end
                end
            end

            xUser = transpose(oldLattice)*oldLattice

            xUser = [xUser[1, 1], xUser[2, 2], xUser[3, 3], xUser[1, 2], xUser[1, 3], xUser[2, 3]]
            diff = xUser

            null = transpose(nullspace(A))

            for i in 1:length(null[1:end, 1])
                diff -= null[i,1:end]*dot(diff,null[i,1:end])/norm(null[i,1:end])
            end

            xCorrected = xUser - diff

            xCorrFull = [xCorrected[1] xCorrected[4] xCorrected[5];xCorrected[4] xCorrected[2] xCorrected[6];xCorrected[5] xCorrected[6] xCorrected[3]]

            newSVD = svd(xCorrFull)

            newSMat = Diagonal(sqrt.(newSVD.S))
            newVtMat = newSVD.Vt

            findUSVD = nothing
            try
                findUSVD = svd(oldLattice*inv(newSMat*newVtMat))
            catch
                continue
            end
            newUMat = findUSVD.U * findUSVD.Vt


            push!(lattices, newUMat*newSMat*newVtMat)
            push!(errors, norm(oldLattice - lattices[end]))
            push!(opCounts, numOp)
        end

        lattice = Array{Float64, 2}(undef, 3, 3)
        numValidOps = 0
        for i in 1:length(errors)+1
            if i == length(errors)+1 || errors[i] > epsilon
                if i == 1
                    i == 2
                end
                lattice = lattices[i-1]
                numValidOps = opCounts[i-1]
                break
            end
        end

        transformation = lattice*latVecInv
        for i in 1:length(atomPositions)
            atomPositions[i] = transformation*atomPositions[i]
        end

        latVecInv = inv(lattice)
        a = lattice[1:end, 1]
        b = lattice[1:end, 2]
        c = lattice[1:end, 3]

        validOperators = []
        for i in 1:numValidOps
            newLattice = transpose(sortTransforms[i]*transpose(lattice))
            push!(validOperators, newLattice*latVecInv)
        end

        return a, b, c, latVecInv, atomPositions, validOperators
    end

    """
        symmSnap(a, b, c, atomPositions, atomTypes[, isDirect = false, fout = "symm_snap.out", epsilon=1e-2])

    Get corrected/reduced lattice vectors, and corrected atom positions to correspond with a Bravais lattice.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `atomPositions::Vector{Vector{Float64}}`: atom positions
    - `atomTypes::_______`:
    - `isDirect::Bool`: indicate if atomPositions are already in direct coordinates
    - `fout::String`: file name to write out additional information to
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function symmSnap(a, b, c, atomPositions, atomTypes, isDirect = false, fout = "symm_snap.out", epsilon=1e-2)
        if isDirect
            atomPositions = toCartesian(a, b, c, atomPositions)
        end

        lattice = Array{Float64, 2}(undef, 3, 3)
        lattice[1:end, 1] = a
        lattice[1:end, 2] = b
        lattice[1:end, 3] = c
        latVecInv = inv(lattice)
        atomPositions = mapAtomsToCell(a, b, c, latVecInv, atomPositions) #new atom positions, but mapped to the unit cell

        lattice = Array{Float64, 2}(undef, 3, 3)
        lattice[1:end, 1] = a
        lattice[1:end, 2] = b
        lattice[1:end, 3] = c
        latVecInv = inv(lattice)

        oldA = deepcopy(a)
        oldB = deepcopy(b)
        oldC = deepcopy(c)

        a, b, c, latVecInv, atomPositions, symmetries = getSymmetryOps(a, b, c, latVecInv, atomPositions)

        output = getLatticeDiff(oldA, a, oldB, b, oldC, c, length(symmetries), [])

        oldAtomPos = deepcopy(atomPositions)

        index = 1
        num_symm = Vector{Int8}(undef, length(atomPositions))

        for num in atomTypes
            a, b, c, atomPositions[index:index+num-1], num_symm[index:index+num-1] = symmetrize(a, b, c, latVecInv, atomPositions[index:index+num-1], symmetries, epsilon)
            index += num
        end

        if isDirect
            atomPositions = toDirect(a, b, c, atomPositions)
        end

        output = getAtomDiff(oldAtomPos, atomPositions, num_symm, output)

        writedlm(fout, output)

        if dot(a, cross(b, c)) < 0
            tempa = a
            a = b
            b = tempa
        end

        return a, b, c, atomPositions
    end

    function snapPoscar(in_file = "POSCAR", out_file = "POSCAR.fixed", log_file = "symm_snap.out", epsilon=1e-2)
        fin=readlines(in_file)
        a=[parse(Float64,i) for i in split(fin[3])]
        b=[parse(Float64,i) for i in split(fin[4])]
        c=[parse(Float64,i) for i in split(fin[5])]
        atom_types=[parse(Int8,i) for i in split(fin[6])]
        is_direct = lowercase(fin[7][1]) == "d" ? true : false
        atom_positions=[[parse(Float64,i) for i in split(fin[j])[1:3]] for j in 8:length(fin)]

        a, b, c, atom_positions = SymmetrySnap.symmSnap(a, b, c, atom_positions, atom_types, is_direct, log_file, epsilon)

        fin[3] = join(a, " ")
        fin[4] = join(b, " ")
        fin[5] = join(c, " ")

        writedlm(out_file, fin)
    end
end
