# filetypes.jl: contains methods needed to parse different file types

"""
    function readXYZ(
        io::IO;
        basis::AbstractVector{AbstractVector{<:Real}},
        spgrp::Integer
        origin::AbstractVector{<:Real}
        prim = false,
        conv = true,
    )

Reads in an XYZ file supplied by the user.

If `warn` is set to `true`, a warning message will be printed to `err` (defaults to `stderr`)
noting that cell vector data is not present in an XYZ file.
"""
function readXYZ(
    io::IO;
    basis::AbstractVector{AbstractVector{<:Real}},
    spgrp::Integer = 0,
    origin::AbstractVector{<:Real} = [0, 0, 0],
    prim = false,
    conv = true,
)
    # Read the file in character by character
    local sites::Vector{SVector{3,Float64}}
    local atoms::Vector{String}
    for (n, ln) in enumerate(eachline(io))
        if n > 2
            # Data starts after the 2nd line
            data = split(ln)
            atoms[n - 2] = data[1]
            sites[n - 2] = SVector{3, Float64}(parse.(Float64, data[2:4]))
        elseif n == 1
            # This is where you get the number of atoms
            natoms = parse(Int, ln)
            atoms = ["" for m = 1:natoms]
            sites = [SVector{3, Float64}(0, 0, 0) for m = 1:natoms]
        end
    end
    # Convert to fractional coordinates if a basis is supplied
    if basis != _basis{3}()
        sites = [inv(hcat(basis...)) * v for v in sites]
    end
    return CrystalStructure{3}(
        basis = basis,
        spgrp = spgrp,
        origin = origin,
        atoms = atoms,
        sites = sites,
        prim = prim,
        conv = conv
    )
end

"""
    readXYZ(filename::AbstractString; warn=false, err=stderr)

Reads in an XYZ file supplied by the user.

If `warn` is set to `true`, a warning message will be printed to `err` (defaults to `stderr`)
noting that cell vector data is not present in an XYZ file.
"""
function readXYZ(
    filename::AbstractString;
    basis::AbstractVector{AbstractVector{<:Real}},
    spgrp::Integer = 0,
    origin::AbstractVector{<:Real} = [0, 0, 0],
    prim = false,
    conv = true,
)
    open(filename) do io
        readXYZ(io, basis=basis, spgrp=spgrp, origin = origin, prim = prim, conv = conv)
    end
end

"""
    writeXYZ(xtal::AbstractCrystal{3}, io::IO)

Writes 3D crystal data to an XYZ file.
"""
function writeXYZ(xtal::AbstractCrystal{3}, io::IO)
    # Write the number of entries
    sz = length(xtal.atoms)
    println(io, string(sz))
    # Write the comment line
    println(io, "Generated by CrystalStructures.jl")
    # Convert basis vectors to a matrix
    hasbasis = (basis == zeros(SMatrix{3,3,Float64}))
    for n in 1:sz
        # Convert the sites from reduced to Cartesian coordinates
        site =  hasbasis ? hcat(xtal.basis...) * xtal.sites[n] : xtal.sites[n]
        @printf(io, "%s\t%f\t%f\t%f\n", xtal.atoms[n], site...)
    end
    return nothing
end

writeXYZ(xtaldata::AbstractCrystalData{3}, io::IO) = writeXYZ(xtaldata.xtal, io)

function writeXYZ(xtal::AbstractCrystal{3}, filename::AbstractString)
    open(filename) do io
        writeXYZ(xtal, io)
    end
end

writeXYZ(xtaldata::AbstractCrystalData{3}, filename::AbstractString) = 
    writeXYZ(xtaldata.xtal, filename)

function writeXYZ(xtaldata::CrystalStructureWithData{3}, filename::AbstractString)
    return writeXYZ(xtaldata.xtal, filename)
end

"""
    readXSF(
        io::IO;
        spgrp::Int = 0,
        origin::Vector{Float64} = [0, 0, 0],
        tryconv=false
    )

Reads in an XCrysDen XSF file from an I/O stream and returns a `CrystalStructure{N}`. Since XSF
data does not contain symmetry information, space group data is not supplied by default, but it may
be included manually by setting the `spgrp` parameter.

This function was written to work around some bugs in an in-house program, `bin2xsf`, that
incorrectly placed the `END_DATAGRID_3D` token on the same line as data. This workaround should not
affect the reading of XSF files with correct syntax.

This function may return an `XSFdata` of arbitrary dimensionality and therefore is not type-stable.
"""
function readXSF(
    io::IO; 
    spgrp::Int = 0,
    origin::AbstractVector{<:Number} = [0, 0, 0],
    tryconv::Bool=false
)
    # Iterator that loops through each line
    iter = eachline(io)
    # Variables needed
    D = 3
    primvec = Vector{Vector{Float64}}()
    convvec = Vector{Vector{Float64}}()
    atomnums = Vector{Int}()
    primcoord = Vector{Vector{Float64}}()
    convcoord = Vector{Vector{Float64}}()
    datagrid_2D = Dict{String, DataGrid{2, Float64}}()
    datagrid_3D = Dict{String, DataGrid{3, Float64}}()
    # Using this while loop because we want to advance the iterator in subloops
    while !isnothing(iter)
        # Get the current line
        # TODO: when this becomes `nothing` it should break, otherwise the assignment fails
        ln = let i = iterate(iter); isnothing(i) ? break : i[1]; end
        # Check for empty lines or comment lines; advance if present
        # Comments are only allowed between blocks in a valid XSF file
        # No need to check within the loops below
        (isempty(ln) || startswith(ln, "#") || all(isspace, ln)) && continue
        # Check for a keyword
        # TODO: This could likely be refactored, perhaps with subfunctions?
        if contains(ln, "DIM-GROUP")
            # Move to the next line
            ln = iterate(iter)[1]
            # Get the dimensionality
            D = parse(Int, split(ln)[1])
        elseif contains(ln, "PRIMVEC")
            # Store vectors here
            primvec = Vector{Vector{Float64}}(undef, D)
            for n in 1:D
                # Move to the next line
                ln = iterate(iter)[1]
                # Get the vector and store it
                primvec[n] = parse.(Float64, split(ln))
            end
        elseif contains(ln, "CONVVEC")
            # Store vectors here
            convvec = Vector{Vector{Float64}}(undef, D)
            for n in 1:D
                # Move to the next line
                ln = iterate(iter)[1]
                # Get the vector and store it
                convvec[n] = parse.(Float64, split(ln))
            end
        elseif contains(ln, "PRIMCOORD")
            # Move to the next line
            ln = iterate(iter)[1]
            # Get number of atoms
            natoms = parse(Int, split(ln)[1])
            # Store atom numbers and vectors here
            atomnums = Vector{Int}(undef, natoms)
            primcoord = Vector{Vector{Float64}}(undef, natoms)
            # Determine the number of atoms
            for n in 1:natoms
                ln = iterate(iter)[1]
                # Every line will be of the following format:
                # [atomname] [x] [y] [z]
                data = split(ln)
                atomnums[n] = parse(Int, data[1])
                primcoord[n] = parse.(Float64, data[2:end])
            end
            # TODO: convert the coordinates in the primitive basis
            if !isempty(primvec)
                primcoord = [inv(hcat(primvec...))*v for v in primcoord]
            end
        elseif contains(ln, "CONVCOORD") && isempty(primcoord)
            # Only do this if primcoord was never given
            # Move to the next line
            ln = iterate(iter)[1]
            # Get number of atoms
            natoms = parse(Int, split(ln)[1])
            # Store atom numbers and vectors here
            atomnums = Vector{Int}(undef, natoms)
            convcoord = Vector{Vector{Float64}}(undef, natoms)
            # Determine the number of atoms
            for n in 1:natoms
                ln = iterate(iter)[1]
                # same format as primcoord
                data = split(ln)
                atomnums[n] = parse(Int, data[1])
                convcoord[n] = parse.(Float64, data[2:end])
            end
        # TODO: implement ATOMS and DATAGRID_2D
        elseif contains.(ln, "BEGIN_BLOCK_DATAGRID_3D") || contains(ln, "BEGIN_BLOCK_DATAGRID3D")
            # Need to use both to avoid issues in bin2xsf
            # Skip the comment line that comes right after
            iterate(iter)
            # Loop until end of block
            while !contains(ln, "END_BLOCK_DATAGRID")
                ln = iterate(iter)[1]
                if contains(ln, "DATAGRID_3D_")
                    # Get the name of the datagrid
                    name = rstrip(ln[findfirst("3D_", ln).stop + 1:end])
                    # Get the dimensionality
                    ln = iterate(iter)[1]
                    dims = parse.(Int, split(ln)[1:3])
                    # Form datagrid as a vector and reshape later
                    datavec = Vector{Float64}(undef, prod(dims))
                    # Get the origin, but I'm not sure if I'll ever need this.
                    ln = iterate(iter)[1]
                    origin = parse.(Float64, split(ln)[1:3])
                    # Check the vectors match primitive or conventional cell
                    vec = Vector{Vector{Float64}}(undef, 3)
                    for n = 1:3
                        ln = iterate(iter)[1]
                        vec[n] = parse.(Float64, split(ln))
                    end
                    if vec != primvec && vec != convvec
                        @warn "Datagrid \"" * name *
                              "\" does not use either primitive or conventional cell vectors."
                    end
                    # Get the actual data
                    entry = 0
                    # Loop until we're either out of entries or we hit an end block
                    while entry < prod(dims) || !contains(ln, "END_DATAGRID")
                        ln = iterate(iter)[1]
                        # This code should tolerate weird spacing of numbers
                        # and work around the bin2xsf bug
                        data = split(ln)
                        # Try to parse and remove unparseable data
                        vals = filter(!isnothing, tryparse.(Float64, data))
                        sz = length(vals)
                        # Add the new data to the datagrid
                        datavec[entry .+ (1:sz)] = vals
                        # Shift the entry counter over
                        entry = entry + sz
                    end
                    # Reshape the datagrid
                    # TODO: probably need logic to determine when to restrict this
                    # But for now, we assume that all XSF files contain 
                    grid = reshape(datavec, dims[1], dims[2], dims[3])[1:end-1, 1:end-1, 1:end-1]
                    # Add it to the dictionary
                    get!(datagrid_3D, name, DataGrid{3,Float64}(vec, grid))
                end
            end
        end
    end
    # Convert the atom numbers to names
    atoms = ELEMENTS[atomnums]
    # Find the right basis
    if isempty(convvec) || !tryconv
        # No conventional vectors found, or specifically requested not to use them
        if isempty(primvec)
            # No cell vectors found
            basis = _basis{D}()
        else
            # Use primitive vectors
            basis = primvec
        end
    else
        basis = convvec
    end
    # Return the CrystalStructure{D}
    return CrystalStructureWithData{D,Float64}(
        CrystalStructure{D}(
            basis = basis,
            spgrp = spgrp,
            # Ignore supplied origin if not 3D
            origin = (D == 3 ? origin : zeros(Float64, D)),
            atoms = atoms,
            sites = (isempty(primcoord) ? convcoord : primcoord),
            conv = basis == convvec,
            prim = basis == primvec,
        ),
        datagrid_3D
    )
end

"""
    readXSF(
        filename::AbstractString; 
        spgrp::Int = 0,
        origin::Vector{Float64} = [0, 0, 0],
        tryconv=true
    )

Reads in an XCrysDen XSF file specified by the user and returns a `CrystalStructure{N}`. Since XSF
data does not contain symmetry information, space group data is not supplied by default, but it may
be included manually by setting the `spgrp` parameter.

This function was written to work around some bugs in an in-house program, `bin2xsf`, that
incorrectly placed the `END_DATAGRID_3D` token on the same line as data. This workaround should not
affect the reading of XSF files with correct syntax.

This function may return an `XSFdata` of arbitrary dimensionality and therefore is not type-stable.
"""
function readXSF(
    filename::AbstractString;
    spgrp::Int = 0,
    origin::AbstractVector{<:Number} = [0, 0, 0],
    tryconv::Bool=false
)
    return open(filename) do io
        readXSF(io, spgrp=spgrp, origin=origin, tryconv=tryconv)
    end
end

"""
    readXSF3D(
        io::IO; 
        spgrp::Int = 0,
        origin::Vector{Float64} = [0, 0, 0],
        tryconv=false
    ) -> CrystalStructure{3}

Reads in an XCrysDen XSF file from an I/O stream and returns a `CrystalStructure{N}`. Since XSF
data does not contain symmetry information, space group data is not supplied by default, but it may
be included manually by setting the `spgrp` parameter.

This function was written to work around some bugs in an in-house program, `bin2xsf`, that
incorrectly placed the `END_DATAGRID_3D` token on the same line as data. This workaround should not
affect the reading of XSF files with correct syntax.

This function is type-stable, but will only work on XSF files representing 3D crystal data (data of
other dimensionality and molecular data are not supported).
"""
function readXSF3D(
    io::IO;
    spgrp::Int = 0,
    origin::AbstractVector{<:Number} = [0, 0, 0],
    tryconv=false
)
    # TODO: probably have some better logic for advancing the iterator
    # Used to get basis vectors
    function getbasis!(itr)
        v = Vector{Vector{Float64}}(undef, 3)
        for n in 1:3
            # Advance to the next line
            ln = iterate(itr)[1]
            v[n] = [parse(Float64, s) for s in split(ln)[1:3]]
        end
        return v
    end

    # Used to get atomic positions
    function getsites!(itr, natoms)
        a = Vector{String}(undef, natoms)
        v = Vector{Vector{Float64}}(undef, natoms)
        for n in 1:natoms
            # Advance to the next line
            ln = iterate(itr)[1]
            # Split the data to get the atom and 3 vector components
            data = split(ln)
            # Conversion to a strong occurs in here!
            a[n] = ELEMENTS[parse(Int, data[1])]
            v[n] = [parse(Float64, s) for s in data[2:4]]
        end
        return (a,v)
    end

    # Used to get datagrids
    function getgrid!(itr, sz; trim = true)
        # Create a new vector with length the same size as the datagrid
        v = Vector{Float64}(undef, prod(sz))
        entry = 0
        ln = ""
        while entry < prod(sz) || !contains(ln, "END_DATAGRID")
            ln = iterate(itr)[1]
            dt = filter(!isnothing, [tryparse(Float64, s) for s in split(ln)])
            v[entry .+ (1:length(dt))] = dt
            # Move the entry counter over
            entry = entry + length(dt)
        end
        # Reshape and trim the edges of the grid, if needed
        return reshape(v, sz[1], sz[2], sz[3])[1:end - trim, 1:end - trim, 1:end - trim]
    end

    # Iterator for the file, reads line by line
    iter = eachline(io)
    # Data needed to construct the 
    primvec = Vector{Vector{Float64}}(undef, 3)
    convvec = Vector{Vector{Float64}}(undef, 3)
    primcoord = Vector{Vector{Float64}}()
    convcoord = Vector{Vector{Float64}}()
    primatoms = Vector{String}()
    convatoms = Vector{String}()
    datagrids = Dict{String,DataGrid{3,Float64}}()
    # The loop
    while !isnothing(iter)
        # Set ln to the next line - break if there is no other line
        ln = let i = iterate(iter); isnothing(i) ? break : i[1]; end
        # Check for empty lines or comment lines; advance if present
        # Comments are only allowed between blocks in a valid XSF file
        # No need to check within the loops below
        (isempty(ln) || startswith(ln, "#") || all(isspace, ln)) && continue
        # Check for keywords
        if contains(ln, "BEGIN_BLOCK_DATAGRID_3D") || contains(ln, "BEGIN_BLOCK_DATAGRID3D")
            # Skip comment line
            iterate(iter)
            # Loop through the whole datagrid block
            while !contains(ln, "END_BLOCK_DATAGRID")
                ln = iterate(iter)[1]
                # Check for a DATAGRID_3D_ keyword
                if contains(ln, "DATAGRID_3D_")
                    # Get the name of the datagrid - this is complicated
                    name = rstrip(ln[findfirst("3D_", ln).stop + 1:end])
                    # Get the grid size
                    ln = iterate(iter)[1]
                    sz = [parse(Int, s) for s in split(ln)]
                    # Get the origin
                    ln = iterate(iter)[1]
                    origin = [parse(Float64, s) for s in split(ln)]
                    # Get the basis vectors for this dataset
                    basis = getbasis!(iter)
                    # Get the actual data
                    grid = getgrid!(iter, sz, trim=true)
                    # Store in the dictionary
                    get!(datagrids, name, DataGrid{3,Float64}(basis, origin, grid))
                end
            end
        elseif contains(ln, "PRIMVEC")
            primvec = getbasis!(iter)
        elseif contains(ln, "CONVVEC")
            convvec = getbasis!(iter)
        elseif contains(ln, "PRIMCOORD")
            ln = iterate(iter)[1]
            natoms = parse(Int, split(ln)[1])
            (primatoms, primcoord) = getsites!(iter, natoms)
        elseif contains(ln, "CONVCOORD")
            ln = iterate(iter)[1]
            natoms = parse(Int, split(ln)[1])
            (convatoms, convcoord) = getsites!(iter, natoms)
        end
    end
    # Choose which basis vectors to use
    # Only try conventional cell data if there are conventional cell vectors/site data
    tryconv = (tryconv && !isempty(convvec) && !isempty(convcoord))
    # Choose which basis vectors to use
    basis = (tryconv ? convvec : primvec)
    # Construct the CrystalStructureWithData{3,Float64}
    return CrystalStructureWithData{3,Float64}(
        CrystalStructure{3}(
            basis = basis,
            spgrp = spgrp,
            origin = origin,
            atoms = tryconv ? convatoms : primatoms,
            sites = [inv(hcat(basis...)) * v for v in (tryconv ? convcoord : primcoord)],
            conv = tryconv,
            prim = (primcoord == convcoord) || !tryconv
        ),
        datagrids
    )
end

function readXSF3D(
    filename::AbstractString;
    spgrp::Int = 0,
    origin::AbstractVector{<:Number} = [0, 0, 0],
    tryconv=false
)
    return open(filename) do io
        readXSF3D(io, spgrp = spgrp, origin = origin, tryconv = tryconv)
    end
end

"""
    function writeXSF3D(
        xtal::CrystalStructureWithData{3},
        io::IO;
        linesize = 4
    )

Writes an XSF file for a 3D crystal structure with data. 
"""
function writeXSF3D(
    xtaldata::CrystalStructureWithData{3, T},
    io::IO;
    linesize = 4
) where T<:Real
    # Needed to format the numbers
    # Matches the format in our in-house bin2xsf
    tostr(x::Number) = rpad(lpad(@sprintf("%.14f", x), 20), 22)
    # Start writing the file
    println(io, "# Written by CrystalStructures.jl\nDIM-GROUP\n 3 1\nPRIMVEC")
    # Write in the primitive cell vectors
    if xtaldata.xtal.prim && !isempty(xtaldata.xtal.basis)
        # Write the primitive cell vectors
        for b in xtaldata.xtal.basis
            println(io, join(tostr.(b)))
        end
    else
        error("Writing crystal structures without primitive vectors is unsupported.")
    end
    # Write PRIMCOORD; skip atoms keyword
    println(io, "PRIMCOORD\n" * lpad(length(xtaldata.xtal.sites), 12) * "  1")
    sites_cart = [hcat(xtaldata.xtal.basis...) * v for v in xtaldata.xtal.sites]
    # TODO: try to figure out what happens when atom names are nonstandard
    atomnums = [ELEMENT_LOOKUP[s] for s in xtaldata.xtal.atoms]
    for (n, v) in zip(atomnums, sites_cart)
        println(io, lpad(string(n), 9) * "  " * join(tostr.(v)))
    end
    # Start writing in the 3D datasets
    println(io, "BEGIN_BLOCK_DATAGRID_3D\nWritten by CrystalStructures.jl")
    for (str,grid) in xtaldata.data
        # Print the datagrid name
        println(io, "    DATAGRID_3D_" * str)
        # Print the datagrid size
        # Note: dimensions of the grid must be increased by 1
        println(io, join([lpad(string(n+1), 6) for n in size(grid.data)]))
        # Print the dataset origin position
        println(io, join([rpad(lpad(@sprintf("%f", x), 20), 22) for x in xtaldata.xtal.origin]))
        # Print the cell vector information for the datagrid
        for b in grid.basis
            println(io, join(tostr.(b)))
        end
        # Print the datagrid
        sz = size(grid.data) .+ 1
        # Counter for inserting newlines
        ctr = 0
        for z = 1:sz[3], y = 1:sz[2], x = 1:sz[1]
            # Generate the new coordinate
            coord = ((x, y, z) .- 1) .% (sz .- 1) .+ 1
            # Write the data at that coordinate
            print(io, tostr(grid.data[coord...]))
            # Insert \n when needed
            ctr = ctr + 1 
            ctr % linesize == 0 && print(io, "\n")
        end
        # Print the end of the datagrid
        println(io, "\n    END_DATAGRID_3D")
    end
    println(io, "END_BLOCK_DATAGRID_3D")
    return nothing
end

function writeXSF3D(
    xtaldata::CrystalStructureWithData{3},
    filename::AbstractString;
    linesize = 4
)
    open(filename, write=true) do io
        writeXSF3D(xtaldata, io, linesize=linesize)
    end
    return nothing
end

# TODO: finish methods for dealing with VASP files
"""
    readPOSCAR(io::IO) -> CrystalStructure{3}

Reads a VASP POSCAR input file, which contains information a crystal's unit cell and atoms. Space
group and origin information can be provided as parameters.
"""
function readPOSCAR(
    io::IO;
    spgrp::Integer = 0,
    origin::AbstractVector{<:Number} = [0, 0, 0],
    prim = true,        # To be expected
    conv = false,
)
    iter = eachline(io)
    # Skip the first line (comment)
    # Second line is a scaling factor
    ln = iterate(iterate(iter))[1]
    scalar = parse(Float64, strip(ln))
    # Next 3 lines are the set of basis vectors
    basis = Vector{Vector{Float64}}(undef, 3)
    for n in 1:3
        ln = iterate(iter)[1]
        basis[n] = parse.(Float64, split(ln))
    end
    ln = iterate(iter)[1]
    # TODO: how to handle the missing atomic symbols?
    # TODO: complete the function
end

readPOSCAR(filename::AbstractString) = open(readPOSCAR, filename)

"""
    readPOSCAR() -> CrystalStructure{3}

Reads any file named `POSCAR` in the current directory as a VASP POSCAR file.
"""
readPOSCAR() = readPOSCAR("POSCAR")