"""
`module  JAC.PhotoExcitation`
    ... a submodel of JAC that contains all methods for computing photo-excitation properties between some initial
        and final-state multiplets.
"""
module PhotoExcitation

    using Printf, ..AngularMomentum, ..Basics,  ..Basics,  ..Defaults, ..ManyElectron, ..Radial, ..PhotoEmission, ..TableStrings

    """
    `struct  PhotoExcitation.Settings`  ... defines a type for the details and parameters of computing photo-excitation  lines.

        + multipoles              ::Array{EmMultipole,1}    ... Specifies the multipoles of the radiation field that are to be included.
        + gauges                  ::Array{UseGauge,1}       ... Specifies the gauges to be included into the computations.
        + calcForStokes           ::Bool                    ... True, if the excitation cross sections are to be calculated (and false otherwise)
                                                                for given Stokes parameter of the incident plane-wave photons.
        + calcPhotonDm            ::Bool                    ... True, if the photon density matrix of a subsequently emitted fluorescence photon
                                                                is to be calculated and false otherwise.
        + calcTensors             ::Bool                    ... True, if statistical tensors of the excited atom are to be calculated, false otherwise.
        + printBefore             ::Bool                    ... True, if all energies and lines are printed before their evaluation.
        + lineSelection           ::LineSelection           ... Specifies the selected levels, if any.
        + photonEnergyShift       ::Float64                 ... An overall energy shift for all photon energies.
        + mimimumPhotonEnergy     ::Float64                 ... minimum transition energy for which (photon) transitions are included into the
                                                                computation.
        + maximumPhotonEnergy     ::Float64                 ... maximum transition energy for which (photon) transitions are included.
        + stokes                  ::ExpStokes               ... Stokes parameters of the incident radiation.
    """
    struct Settings
        multipoles                ::Array{EmMultipole,1}
        gauges                    ::Array{UseGauge,1}
        calcForStokes             ::Bool
        calcPhotonDm              ::Bool
        calcTensors               ::Bool
        printBefore               ::Bool
        lineSelection             ::LineSelection
        photonEnergyShift         ::Float64
        mimimumPhotonEnergy       ::Float64
        maximumPhotonEnergy       ::Float64
        stokes                    ::ExpStokes
    end

    """
    `PhotoExcitation.Settings()`  ... 'empty' constructor for the default values of photo-excitation line computations
    """
    function Settings()
        Settings(EmMultipole[], UseGauge[], false, false, false, false, LineSelection(), 0., 0., 0., Basics.ExpStokes())
    end

    """
    `PhotoExcitation.Settings(set::PhotoExcitation.Settings;`

            multipoles=..,          gauges=..,                  calcForStokes=..,           calcPhotonDm=..,
            calcTensors=..,         printBefore=..,             lineSelection=..,
            photonEnergyShift=..,   mimimumPhotonEnergy=..,     maximumPhotonEnergy=..,     stokes=..)

        ... constructor for modifying the given PhotoExcitation.Settings by 'overwriting' the previously selected parameters.
    """
    function Settings(set::PhotoExcitation.Settings;
        multipoles::Union{Nothing,Array{EmMultipole,1}}=nothing,        gauges::Union{Nothing,Array{UseGauge,1}}=nothing,
        calcForStokes::Union{Nothing,Bool}=nothing,                     calcPhotonDm::Union{Nothing,Bool}=nothing,
        calcTensors::Union{Nothing,Bool}=nothing,                       printBefore::Union{Nothing,Bool}=nothing,
        lineSelection::Union{Nothing,LineSelection}=nothing,
        photonEnergyShift::Union{Nothing,Float64}=nothing,              mimimumPhotonEnergy::Union{Nothing,Float64}=nothing,
        maximumPhotonEnergy::Union{Nothing,Float64}=nothing,            stokes::Union{Nothing,ExpStokes}=nothing)

        if  multipoles          == nothing   multipolesx          = set.multipoles              else  multipolesx          = multipoles            end
        if  gauges              == nothing   gaugesx              = set.gauges                  else  gaugesx              = gauges                end
        if  calcForStokes       == nothing   calcForStokesx       = set.calcForStokes           else  calcForStokesx       = calcForStokes         end
        if  calcPhotonDm        == nothing   calcPhotonDmx        = set.calcPhotonDm            else  calcPhotonDmx        = calcPhotonDm          end
        if  calcTensors         == nothing   calcTensorsx         = set.calcTensors             else  calcTensorsx         = calcTensors           end
        if  printBefore         == nothing   printBeforex         = set.printBefore             else  printBeforex         = printBefore           end
        if  lineSelection       == nothing   lineSelectionx       = set.lineSelection           else  lineSelectionx       = lineSelection         end
        if  photonEnergyShift   == nothing   photonEnergyShiftx   = set.photonEnergyShift       else  photonEnergyShiftx   = photonEnergyShift     end
        if  mimimumPhotonEnergy == nothing   mimimumPhotonEnergyx = set.mimimumPhotonEnergy     else  mimimumPhotonEnergyx = mimimumPhotonEnergy   end
        if  maximumPhotonEnergy == nothing   maximumPhotonEnergyx = set.maximumPhotonEnergy     else  maximumPhotonEnergyx = maximumPhotonEnergy   end
        if  stokes              == nothing   stokesx              = set.stokes                  else  stokesx              = stokes                end

        Settings( multipolesx, gaugesx, calcForStokesx, calcPhotonDmx, calcTensorsx, printBeforex, lineSelectionx,
                  photonEnergyShiftx, mimimumPhotonEnergyx, maximumPhotonEnergyx, stokesx)
    end

    # `Base.show(io::IO, settings::PhotoExcitation.Settings)`  ... prepares a proper printout of the variable settings::PhotoExcitation.Settings.
    function Base.show(io::IO, settings::PhotoExcitation.Settings)
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "use-gauges:               $(settings.gauges)  ")
        println(io, "calcForStokes:            $(settings.calcForStokes)  ")
        println(io, "calcPhotonDm:             $(settings.calcPhotonDm)  ")
        println(io, "calcTensors:              $(settings.calcTensors)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "lineSelection:            $(settings.lineSelection)  ")
        println(io, "photonEnergyShift:        $(settings.photonEnergyShift)  ")
        println(io, "mimimumPhotonEnergy:      $(settings.mimimumPhotonEnergy)  ")
        println(io, "maximumPhotonEnergy:      $(settings.maximumPhotonEnergy)  ")
        println(io, "stokes:                   $(settings.stokes)  ")
    end

    """
    `struct  PhotoExcitation.Line`
        ... defines a type for a photo-excitation line that may include the definition of sublines and their
            corresponding amplitudes.

        + initialLevel   ::Level                       ... initial-(state) level
        + finalLevel     ::Level                       ... final-(state) level
        + omega          ::Float64                     ... Transition frequency of this line; can be shifted w.r.t. the level energies.
        + oscStrength    ::EmProperty                  ... Absorption oscillator strength
        + crossSection   ::EmProperty                  ... Total cross section of this line.
        + staTensor      ::Array{TensorComp,1}         ... Array of statistical tensor components rho_kq
        + hasSublines    ::Bool                        ... Determines whether the individual sublines are defined in terms of their
                                                           multipolarity, amplitude, or not; cf. PhotoEmission.Channel
        + channels       ::Array{PhotoEmission.Channel,1}  ... List of radiative (photon) channels
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        omega            ::Float64
        oscStrength      ::EmProperty
        crossSection     ::EmProperty
        staTensor        ::Array{TensorComp,1}
        hasSublines      ::Bool
        channels         ::Array{PhotoEmission.Channel,1}
    end

    """
    `PhotoExcitation.Line(initialLevel::Level, finalLevel::Level, omega::Float64, crossSection::EmProperty)`
        ... constructor an photo-excitation line between a specified initial and final level.
    """
    function Line(initialLevel::Level, finalLevel::Level, omega::Float64, crossSection::EmProperty)
        Line(initialLevel, finalLevel, omega, EmProperty(0., 0.), crossSection, false, PhotoEmission.Channel[])
    end

    # `Base.show(io::IO, line::PhotoExcitation.Line)`  ... prepares a proper printout of the variable line::PhotoExcitation.Line.
    function Base.show(io::IO, line::PhotoExcitation.Line)
        println(io, "initialLevel:         $(line.initialLevel)  ")
        println(io, "finalLevel:           $(line.finalLevel)  ")
        println(io, "omega:                $(line.omega)  ")
        println(io, "oscStrength:          $(line.oscStrength)  ")
        println(io, "crossSection:         $(line.crossSection)  ")
        println(io, "staTensor:            $(line.staTensor)  ")
        println(io, "hasSublines:          $(line.hasSublines)  ")
        println(io, "channels:             $(line.channels)  ")
    end

    """
    `PhotoExcitation.computeAmplitudesProperties(line::PhotoExcitation.Line, grid::Radial.Grid, settings::PhotoExcitation.Settings; printout::Bool=true)`
        ... to compute all amplitudes and properties of the given line; a line::PhotoExcitation.Line is returned for which
            the amplitudes and properties have now been evaluated.
    """
    function  computeAmplitudesProperties(line::PhotoExcitation.Line, grid::Radial.Grid, settings::PhotoExcitation.Settings; printout::Bool=true)
        newChannels = PhotoEmission.Channel[];    Ji2 = Basics.twice(line.initialLevel.J);    Jf2 = Basics.twice(line.finalLevel.J)
        for  channel  in  line.channels
            amplitude = PhotoEmission.amplitude("absorption", channel.multipole, channel.gauge, line.omega,
                                                line.finalLevel, line.initialLevel, grid, printout=printout)
            push!( newChannels, PhotoEmission.Channel(channel.multipole, channel.gauge, amplitude) )
        end
        # Calculate the absorption oscillator strength
        oscCoulomb = oscBabushkin = 0.;   omega = line.omega;   alpha = Defaults.getDefaults("alpha")
        for  channel  in  newChannels
            ##x wa = line.omega / (Basics.twice(channel.multipole.L) + 1) * (alpha * line.omega)^(2*channel.multipole.L - 2)
            wa = 8pi * alpha * line.omega / (Ji2 + 1) * (Jf2 + 1)
            if      channel.gauge == Basics.Coulomb     oscCoulomb   = oscCoulomb    +  channel.amplitude * conj(channel.amplitude) * wa
            elseif  channel.gauge == Basics.Babushkin   oscBabushkin = oscBabushkin  +  channel.amplitude * conj(channel.amplitude) * wa
            elseif  channel.gauge == Basics.Magnetic    oscBabushkin = oscBabushkin  +  channel.amplitude * conj(channel.amplitude) * wa
                                                        oscCoulomb   = oscCoulomb    +  channel.amplitude * conj(channel.amplitude) * wa
            else    error("stop a")
            end
        end
        # Calculate the crossSection and stastistical tensors if requested
        csCoulomb = csBabushkin = 0.;
        for  channel  in  newChannels
            if      channel.gauge == Basics.Coulomb     csCoulomb   = csCoulomb    +  channel.amplitude * conj(channel.amplitude)
            elseif  channel.gauge == Basics.Babushkin   csBabushkin = csBabushkin  +  channel.amplitude * conj(channel.amplitude)
            elseif  channel.gauge == Basics.Magnetic    csBabushkin = csBabushkin  +  channel.amplitude * conj(channel.amplitude)
                                                        csCoulomb   = csCoulomb    +  channel.amplitude * conj(channel.amplitude)
            else    error("stop b")
            end
        end
        oscFactor    = (Ji2 + 1) / (Jf2 + 1) / 2. * Defaults.getDefaults("speed of light: c")^3 / line.omega^2
        csFactor     = 2 * pi^3 * Defaults.getDefaults("alpha") / (line.omega * (Ji2 + 1))
        oscStrength  = EmProperty(oscFactor*oscCoulomb, oscFactor*oscBabushkin)
        crossSection = EmProperty(csFactor*csCoulomb, csFactor*csBabushkin)
        staTensor    = TensorComp[]
        if  settings.calcTensors    push!(staTensor, TensorComp(0, 0, 1.))      end

        newLine = PhotoExcitation.Line( line.initialLevel, line.finalLevel, line.omega, oscStrength, crossSection, staTensor, true, newChannels)
        return( newLine )
    end

    """
    `PhotoExcitation.computeCrossSection(line::PhotoExcitation.Line, stokes::ExpStokes)`
        ... to compute the excitation cross section for the excitation of unpolarized atoms by plane-wave photons, whose polarization
            is described by the given (experimental) Stokes parameters. A cs::EmProperty is returned.
    """
    function  computeCrossSection(line::PhotoExcitation.Line, stokes::ExpStokes)
        csCoulomb = -2.;    csBabushkin = -2.
        cs = EmProperty( csCoulomb, csBabushkin )
        return( cs )
    end

    """
    `PhotoExcitation.computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                      settings::PhotoExcitation.Settings; output=true)`
        ... to compute the photo-excitation amplitudes and all properties as requested by the given settings. A list
            of lines::Array{PhotoExcitation.Lines} is returned.
    """
    function  computeLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                           settings::PhotoExcitation.Settings; output=true)
        println("")
        printstyled("PhotoExcitation.computeLines(): The computation of the excitation cross sections, etc. starts now ... \n", color=:light_green)
        printstyled("----------------------------------------------------------------------------------------------------- \n", color=:light_green)
        println("")
        lines = PhotoExcitation.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        if  settings.printBefore    PhotoExcitation.displayLines(lines)    end
        # Calculate all amplitudes and requested properties
        newLines = PhotoExcitation.Line[]
        for  line in lines
            newLine = PhotoExcitation.computeAmplitudesProperties(line, grid, settings)
            push!( newLines, newLine)
        end
        # Print all results to screen
        PhotoExcitation.displayCrossSections(stdout, newLines, settings)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary    PhotoExcitation.displayCrossSections(iostream, newLines, settings)   end
        #
        if    output    return( lines )
        else            return( nothing )
        end
    end

    """
    `PhotoExcitation.computeLinesCascade(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                         settings::PhotoExcitation.Settings; output::Bool=true, printout::Bool=true)`
        ... to compute the excitation (absorption) transition amplitudes and all properties as requested by the given settings.
            The computations and printout is adapted for larger cascade computations by including only lines with at least one channel
            and by sending all printout to a summary file only. A list of lines::Array{PhotoExcitation.Lines} is returned.
    """
    function  computeLinesCascade(finalMultiplet::Multiplet, initialMultiplet::Multiplet, grid::Radial.Grid,
                                         settings::PhotoExcitation.Settings; output::Bool=true, printout::Bool=true)
        # Define a common subshell list for both multiplets
        subshellList = Basics.generate("subshells: ordered list for two bases", finalMultiplet.levels[1].basis, initialMultiplet.levels[1].basis)
        Defaults.setDefaults("relativistic subshell list", subshellList; printout=false)
        lines = PhotoExcitation.determineLines(finalMultiplet, initialMultiplet, settings)
        # Display all selected lines before the computations start
        # if  settings.printBefore    PhotoExcitation.displayLines(lines)    end
        # Calculate all amplitudes and requested properties
        newLines = PhotoExcitation.Line[]
        for  (i,line)  in  enumerate(lines)
            if  rem(i,50) == 0    println("> Excitation line $i:")   end
            newLine = PhotoExcitation.computeAmplitudesProperties(line, grid, settings, printout=printout)
            #
            # Don't add this line if it does not contribute to the decay
            wa = 0.
            for  ch in newLine.channels   wa = wa + abs(ch.amplitude)^2    end
            if   wa == 0.    continue    end
            push!( newLines, newLine)
        end
        # Print all results to a summary file, if requested
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary   PhotoExcitation.displayCrossSections(iostream, newLines, settings)    end
        #
        if    output    return( newLines )
        else            return( nothing )
        end
    end

    """
    `PhotoExcitation.computeStatisticalTensor(k::Int64, q::Int64, line::PhotoExcitation.Line, stokes::ExpStokes)`
        ... to compute the statistical tensor (component) rho_{k,q} of the final level for the excitation of unpolarized atoms by
            plane-wave photons, whose polarization is described by the given (experimental) Stokes parameters.
            A rho_kq::EmProperty is returned.
    """
    function  computeStatisticalTensor(k::Int64, q::Int64, line::PhotoExcitation.Line, stokes::ExpStokes)
        if  k == q  ||  k == -q   rhoCoulomb = -3.;    rhoBabushkin = -3.    else   rhoCoulomb = 0.;    rhoBabushkin = 0.    end
        rho = EmProperty( rhoCoulomb, rhoBabushkin )
        return( rho )
    end

    """
    `PhotoExcitation.determineChannels(finalLevel::Level, initialLevel::Level, settings::PhotoExcitation.Settings)`
        ... to determine a list of PhotoExcitation.Channel for a transitions from the initial to final level and by taking
            into account the particular settings of for this computation; an Array{PhotoExcitation.Channel,1} is returned.
    """
    function determineChannels(finalLevel::Level, initialLevel::Level, settings::PhotoExcitation.Settings)
        channels = PhotoEmission.Channel[];
        symi = LevelSymmetry(initialLevel.J, initialLevel.parity);    symf = LevelSymmetry(finalLevel.J, finalLevel.parity)
        for  mp in settings.multipoles
            if   AngularMomentum.isAllowedMultipole(symi, mp, symf)
                hasMagnetic = false
                for  gauge in settings.gauges
                    # Include further restrictions if appropriate
                    if     string(mp)[1] == 'E'  &&   gauge == Basics.UseCoulomb      push!(channels, PhotoEmission.Channel(mp, Basics.Coulomb,   0.) )
                    elseif string(mp)[1] == 'E'  &&   gauge == Basics.UseBabushkin    push!(channels, PhotoEmission.Channel(mp, Basics.Babushkin, 0.) )
                    elseif string(mp)[1] == 'M'  &&   !(hasMagnetic)               push!(channels, PhotoEmission.Channel(mp, Basics.Magnetic,  0.) );
                                                        hasMagnetic = true;
                    end
                end
            end
        end
        ##x println("PhotoExcitation.determineChannels-aa: channels = $channels ")
        return( channels )
    end

    """
    `PhotoExcitation.determineLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, settings::PhotoExcitation.Settings)`
        ... to determine a list of photo-excitation Line's for transitions between the levels from the given initial- and
            final-state multiplets and by taking into account the particular selections and settings for this computation;
            an Array{PhotoExcitation.Line,1} is returned. Apart from the level specification, all physical properties are set to
            zero during the initialization process.
    """
    function  determineLines(finalMultiplet::Multiplet, initialMultiplet::Multiplet, settings::PhotoExcitation.Settings)
        lines = PhotoExcitation.Line[]
        for  iLevel  in  initialMultiplet.levels
            for  fLevel  in  finalMultiplet.levels
                if  Basics.selectLevelPair(iLevel, fLevel, settings.lineSelection)
                    omega    = fLevel.energy - iLevel.energy + settings.photonEnergyShift
                    if   omega == 0.  ||  omega < settings.mimimumPhotonEnergy  ||  omega > settings.maximumPhotonEnergy    continue   end
                    channels = PhotoExcitation.determineChannels(fLevel, iLevel, settings)
                    if   length(channels) == 0   continue   end
                    push!( lines, PhotoExcitation.Line(iLevel, fLevel, omega, EmProperty(0., 0.), EmProperty(0., 0.), TensorComp[], true, channels) )
                end
            end
        end
        return( lines )
    end

    """
    `PhotoExcitation.displayLines(lines::Array{PhotoExcitation.Line,1})`
        ... to display a list of lines and channels that have been selected due to the prior settings. A neat table of all
            selected transitions and energies is printed but nothing is returned otherwise.
    """
    function  displayLines(lines::Array{PhotoExcitation.Line,1})
        nx = 120
        println(" ")
        println("  Selected photo-excitation lines:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        sa = sa * TableStrings.flushleft(30, "List of multipoles"; na=4);             sb = sb * TableStrings.hBlank(34)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                         fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
            mpGaugeList = Tuple{Basics.EmMultipole,Basics.EmGauge}[]
            for  i in 1:length(line.channels)
                push!( mpGaugeList, (line.channels[i].multipole, line.channels[i].gauge) )
            end
            ##x println("PhotoExcitation.diplayLines-ad: mpGaugeList = ", mpGaugeList)
            sa = sa * TableStrings.multipoleGaugeTupels(50, mpGaugeList)
            println( sa )
        end
        println("  ", TableStrings.hLine(nx))
        #
        return( nothing )
    end

    """
    `PhotoExcitation.displayCrossSections(stream::IO, lines::Array{PhotoExcitation.Line,1}, settings::PhotoExcitation.Settings)`
        ... to list all results, energies, cross sections, etc. of the selected lines. A neat table is printed but nothing
            is returned otherwise.
    """
    function  displayCrossSections(stream::IO, lines::Array{PhotoExcitation.Line,1}, settings::PhotoExcitation.Settings)
        #
        # Photoexcitation resonance strength as derived from oscillator strength
        nx = 122
        println(stream, " ")
        println(stream, "  Photoexcitation resonance strength as derived from oscillator strength:")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"   ; na=3);
        sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
        sa = sa * TableStrings.center(10, "Multipoles"; na=3);                        sb = sb * TableStrings.hBlank(13)
        sa = sa * TableStrings.center(24, "Cou -- f_ik -- Bab"; na=1);
        sb = sb * TableStrings.center(24, "     absorption   "; na=2)
        sa = sa * TableStrings.center(24, "Cou -- S -- Bab"; na=3);
        sb = sb * TableStrings.center(24, "   [Mb eV]     "; na=3)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                         fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
            multipoles = EmMultipole[]
            for  ch in line.channels
                multipoles = push!( multipoles, ch.multipole)
            end
            multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
            sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=3)
            sa = sa * @sprintf("%.4e", line.oscStrength.Coulomb)   * "   "
            sa = sa * @sprintf("%.4e", line.oscStrength.Babushkin) * "    "
            sa = sa * @sprintf("%.4e", line.oscStrength.Coulomb   * 109.761)   * "   "
            sa = sa * @sprintf("%.4e", line.oscStrength.Babushkin * 109.761)   * "    "
            println(stream, sa)
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        #
        # Photoexcitation cross sections from line strength for (completely) linearly-polarized plane-wave photons
        nx = 103
        println(stream, " ")
        println(stream, "  Photoexcitation cross sections for (completely) linearly-polarized plane-wave photons:  ... not yet correct !!")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
        sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
        sa = sa * TableStrings.center(14, "Energy"   ; na=3);
        sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
        sa = sa * TableStrings.center(10, "Multipoles"; na=2);                        sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(28, "Cou -- cross section -- Bab"; na=0);
        sb = sb * TableStrings.center(28, TableStrings.inUnits("cross section")*"       "*
                                          TableStrings.inUnits("cross section"); na=1)
        println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
        #
        for  line in lines
            sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                         fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
            sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
            sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
            multipoles = EmMultipole[]
            for  ch in line.channels
                multipoles = push!( multipoles, ch.multipole)
            end
            multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
            sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=5)
            sa = sa * @sprintf("%.4e", Defaults.convertUnits("cross section: from atomic", line.crossSection.Coulomb))     * "   "
            sa = sa * @sprintf("%.4e", Defaults.convertUnits("cross section: from atomic", line.crossSection.Babushkin))   * "    "
            ##x sa = sa * @sprintf("%.6e", line.anisotropy.Coulomb)      * "    "
            ##x sa = sa * @sprintf("%.6e", line.anisotropy.Babushkin)    * "    "
            println(stream, sa)
        end
        println(stream, "  ", TableStrings.hLine(nx))
        #
        #
        if  settings.calcForStokes
            stokes = settings.stokes
            println(stream, " ")
            println(stream, "  Photoexcitation cross sections for incident plane-wave photons with given $stokes:  ... not yet correct !!")
            println(stream, " ")
            println(stream, "  ", TableStrings.hLine(nx))
            sa = "  ";   sb = "  "
            sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
            sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
            sa = sa * TableStrings.center(14, "Energy"   ; na=3);
            sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
            sa = sa * TableStrings.center(10, "Multipoles"; na=3);                        sb = sb * TableStrings.hBlank(14)
            sa = sa * TableStrings.center(30, "Cou -- Cross section -- Bab"; na=2);
            sb = sb * TableStrings.center(30, TableStrings.inUnits("cross section")*"          "*
                                                  TableStrings.inUnits("cross section"); na=2)
            println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
            for  line in lines
                sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                             fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
                sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
                sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
                sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
                multipoles = EmMultipole[]
                for  ch in line.channels
                    multipoles = push!( multipoles, ch.multipole)
                end
                multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
                sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=3)
                crossSection = PhotoExcitation.computeCrossSection(line, stokes)
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("cross section: from atomic", crossSection.Coulomb))     * "    "
                sa = sa * @sprintf("%.6e", Defaults.convertUnits("cross section: from atomic", crossSection.Babushkin))   * "    "
                println(stream, sa)
            end
            println(stream, "  ", TableStrings.hLine(nx))
        end
        #
        #
        if  settings.calcTensors
            nx = 145
            stokes = settings.stokes
            println(stream, " ")
            println(stream, "  Statistical tensors rho_kq  and alignment parameters A_kq for the excitation by incident plane-wave photons")
            println(stream, "  with given $stokes:  ... not yet correct !!")
            println(stream, " ")
            println(stream, "  ", TableStrings.hLine(nx))
            sa = "  ";   sb = "  "
            sa = sa * TableStrings.center(18, "i-level-f"; na=0);                         sb = sb * TableStrings.hBlank(18)
            sa = sa * TableStrings.center(18, "i--J^P--f"; na=4);                         sb = sb * TableStrings.hBlank(22)
            sa = sa * TableStrings.center(14, "Energy"   ; na=3);
            sb = sb * TableStrings.center(14,TableStrings.inUnits("energy"); na=3)
            sa = sa * TableStrings.center(10, "Multipoles"; na=1);                        sb = sb * TableStrings.hBlank(12)
            sa = sa * TableStrings.center(12, "     k    q    "; na=0);                   sb = sb * TableStrings.hBlank(14)
            sa = sa * TableStrings.center(26, "Cou -- rho_kq -- Bab"; na=5);              sb = sb * TableStrings.hBlank(14)
            sa = sa * TableStrings.center(26, "Cou --  A_kq  -- Bab"; na=2);              sb = sb * TableStrings.hBlank(14)
            println(stream, sa);    println(stream, sb);    println(stream, "  ", TableStrings.hLine(nx))
            for  line in lines
                sa  = "";    isym = LevelSymmetry( line.initialLevel.J, line.initialLevel.parity)
                             fsym = LevelSymmetry( line.finalLevel.J,   line.finalLevel.parity)
                sa = sa * TableStrings.center(18, TableStrings.levels_if(line.initialLevel.index, line.finalLevel.index); na=2)
                sa = sa * TableStrings.center(18, TableStrings.symmetries_if(isym, fsym); na=4)
                sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", line.omega)) * "    "
                multipoles = EmMultipole[]
                for  ch in line.channels
                    multipoles = push!( multipoles, ch.multipole)
                end
                multipoles = unique(multipoles);   mpString = TableStrings.multipoleList(multipoles) * "          "
                sa = sa * TableStrings.flushleft(11, mpString[1:10];  na=3)
                println(stream, sa);   sc = TableStrings.hBlank( length(sa)-1 )
                tc00   = PhotoExcitation.computeStatisticalTensor(0, 0, line, stokes)
                for  k = 2:2
                    for  q = -k:k
                        tc = PhotoExcitation.computeStatisticalTensor(k, q, line, stokes)
                        if  abs(tc.Coulomb) == abs(tc.Babushkin) == 0.    continue    end
                        #
                        aCoulomb = tc.Coulomb / tc00.Coulomb;    aBabushkin = tc.Babushkin / tc00.Babushkin
                        sd = sc * TableStrings.level(k) * TableStrings.level(q) * "   "
                        sd = sd * @sprintf("%.6e", tc.Coulomb)     * "  "
                        sd = sd * @sprintf("%.6e", tc.Babushkin)   * "    "
                        sd = sd * @sprintf("%.6e", aCoulomb)       * "  "
                        sd = sd * @sprintf("%.6e", aBabushkin)     * "    "
                        println(stream, sd)
                    end
                end
            end
            println(stream, "  ", TableStrings.hLine(nx))
        end
        #
        #
        if  settings.calcTensors   &&   settings.calcPhotonDm
            stokes = settings.stokes
            println(stream, " ")
            println(stream, "  Reduced density matrix of the fluorescence photons at selected solid angles for the prior excitation by incident plane-wave photons")
            println(stream, "  with given $stokes:")
            println(stream, " ")
            println(stream, "  Not yet implemented !!")
        end
        #
        return( nothing )
    end

end # module


using Images # requires Package Images.jl
using Colors # requires Package Colors.jl
using Mocha
using Compat

type ImageClassifier
  net           :: Net
  channel_order :: NTuple{3,Int} # The channel order of the trained net, (1,2,3) means RGB
  sp_order      :: NTuple{2,Int} # The spatial order (1,2) means width-height (row major)
  grayscale     :: Bool          # Whether we are working with gray (single-channel) images
  classes       :: Vector        # Names for each class

  data_layer :: MemoryDataLayer
  pred_blob  :: Blob
  pred       :: Array
  data_type  :: Type
  batch_size :: Int
  image_wh   :: NTuple{2, Int}
end

function ImageClassifier(net::Net, softmax_blob_name::Symbol;
    channel_order = (1,2,3), # default RGB
    sp_order  = (1,2),       # default row-major (width,height,channels)
    grayscale = false,       # default not grayscale
    classes = [],            # default no name information available
    )

  if length(net.data_layers) != 1 || !isa(net.layers[net.data_layers[1]], MemoryDataLayer)
    error("The network should contain exactly one MemoryDataLayer")
  end
  data_layer = net.layers[net.data_layers[1]]
  batch_size = data_layer.batch_size
  data_type  = eltype(data_layer.data[1])
  image_wh   = (size(data_layer.data[1],1), size(data_layer.data[1],2))

  @assert length(data_layer.data) == 1
  if size(data_layer.data[1], 4) != batch_size
    # to make our life easier
    error("Please make the num dimension ($(size(data_layer.data[1],4))) of MemoryDataLayer's data source the same to batch_size ($(batch_size))")
  end

  pred_blob = net.output_blobs[softmax_blob_name]
  pred = Array(data_type, size(pred_blob))

  return ImageClassifier(net, channel_order, sp_order, grayscale, classes,
      data_layer, pred_blob, pred, data_type, batch_size, image_wh)
end

################################################################################
# Interfaces
################################################################################
#-- Classify a single image via filename
function classify(classifier::ImageClassifier, filename::AbstractString)
  results = classify(classifier, AbstractString[filename])
  return results[1]
end

#-- Classify a bunch of images via filename
function classify{T<:AbstractString}(classifier::ImageClassifier, filenames::Vector{T})
  return classify(classifier, map(Images.imread, filenames))
end

#-- Classify a single image
function classify(classifier::ImageClassifier, image::Image)
  results = classify(classifier, Image[image])
  return results[1]
end

#-- Classify a bunch of images
function classify(classifier::ImageClassifier, images::Vector{Image})
  results = classify_batch(classifier, images, 1)
  idx = classifier.batch_size+1
  while idx <= length(images)
    ret_batch = classify_batch(classifier, images, idx)
    results = [results, ret_batch]
    idx += length(res_batch)
  end

  return results
end

################################################################################
# Implementations
################################################################################
#-- Propress images
function preprocess(classifier::ImageClassifier, images::Vector{Image})
  map(images) do image
    if (width(img), height(img)) != classifier.image_wh
      error("Image size ($(width(img)),$(height(img))) does not match the network input $(classifier.image_wh), and Julia does not have imresize yet........")
    end

    if classifier.grayscale
      if colorspace(image) != "Gray"
        image = convert(Image{Gray}, image)
      end
    else
      image = convert(Image{RGB}, image)
      image = separate(image) # separate color channels
    end

    data = convert(Array, image)
    data = reshape(data, size(data,1),size(data,2),size(data,3))
    if spatialorder(image) == ["x","y"]
      # row major
    elseif spatialorder(image) == ["y","x"]
      data = permutedims(data, (2,1,3)) # permute to row-major
    else
      error("Unknown spatialorder: $(spatialorder(image))")
    end

    # now data is in canonical row-major RGB/Gray format
    if classifier.sp_order != (1,2)
      @assert classifier.sp_order == (2,1)
      data = permutedims(data, (2,1,3))
    end

    if classifier.grayscale
      data2 = convert(Array{classifier.data_type}, data)
    else
      data2 = Array(classifier.data_type, tuple(classifier.image_wh..., 3))
      for i = 1:3
        data2[:,:,i] = convert(Array{classifier.data_type},
            data[:,:,classifier.channel_order[i]])
      end
    end

    data2 # pre-processed raw image data
  end
end

#-- Classify a batch or less
function classify_batch(classifier::ImageClassifier, images::Vector{Image}, idx::Int)
  #-- prepare the data
  idx_end = min(idx+classifier.batch_size-1, length(images))
  images = preprocess(classifier, images[idx:idx_end])
  for i = 1:length(images)
    classifier.data_layer.data[1][:,:,:,i] = images[i]
  end

  #-- run the network
  forward(classifier.net)
  copy!(classifier.pred, classifier.pred_blob)
  results = map(1:length(images)) do i
    pred = classifier.pred[:,i]
    i_class = indmax(pred)
    if !isempty(classifier.classes)
      ret = classifier.classes[i_class]
    else
      ret = string(i_class-1)
    end
    return (pred, ret)
  end

  return results
end


const keywordlist = [
    "ELEMENT",
    "FUNCTION",
    "FUNCT",
    "TYPE_DEFINITION",
    "PHASE",
    "CONSTITUENT",
    "CONST",
    "PARAMETER",
    "PAR",
]

function read_tdb(io::IO)
    elems = Vector{Element}()
    funcs = Vector{GFunction}()
    phass = Vector{Phase}()
    types = Vector{TypeDefinition}()

    # Parse each elements
    while !eof(io)
        block = readline(io, keep = true)

        block = replace(block, '\t' => "") # remove tabs

        isempty(block) && continue # blank line
        !occursin(r"\S+", block) && continue # blank line
        occursin(r"\$", block) && continue # comment line

        i = findfirst('!', block)
        if isnothing(i)
            block *= readuntil(io, '!')
        else
            block = block[1:i-1] # remove '!'
        end

        block = replace(block, '\r' => "")
        block = replace(block, '\n' => ' ')
        strs = split(block)
        isempty(strs) && continue

        keyword = uppercase(strs[1])

        !(keyword in keywordlist) && continue
        try
            if keyword == "ELEMENT"
                elem = parse_element(block)
                push!(elems, elem)
            elseif keyword in ["FUNCTION", "FUNCT"]
                func = parse_gfunction(block)
                push!(funcs, func)
            elseif keyword == "TYPE_DEFINITION"
                type = parse_typedefinition(block)
                push!(types, type)
            elseif keyword == "PHASE"
                phas = parse_phase(block)
                push!(phass, phas)
            elseif keyword in ["CONSTITUENT", "CONST"]
                phasname, cons = parse_constituent(block)
                phas = filter(phas -> uppercase(phas.name) == uppercase(phasname), phass)
                @assert length(phas) == 1
                phas[1].cons = cons
            elseif keyword in ["PARAMETER", "PAR"]
                phasname, param = parse_parameter(block)
                isnothing(param) && continue
                phas = filter(phas -> uppercase(phas.name) == uppercase(phasname), phass)
                @assert length(phas) == 1
                push!(phas[1].params, param)
            end
        catch
            error("Unrecognized description in the tdb file:\n$block")
        end
    end

    return Database(elems, funcs, phass, types)
end

function read_tdb(tdb::AbstractString)
    open(tdb, "r") do io
        read_tdb(io)
    end
end

function parse_element(text::AbstractString)
    strs = split(text)
    @assert length(strs) == 6
    @assert match(r"ELEMENT"i, strs[1]) ≠ ""
    name = format_element(strs[2])
    refstate = strs[3]
    mass, H298, S298 = parse.(Float64, strs[4:6])
    return Element(name, refstate, mass, H298, S298)
end

function parse_gfunction(text::AbstractString)
    strs = split(text)
    @assert match(r"FUNCTION"i, strs[1]) ≠ ""
    name = strs[2]
    temp, funcstr = parse_function(name, strs[3:end])
    return GFunction(name, temp, funcstr)
end

function parse_function(name::AbstractString, strs::Vector{SubString{String}})
    Ts = Vector{AbstractString}()
    funcstrs = Vector{AbstractString}()

    push!(Ts, strs[1]) # ex. 298.15
    k = 2

    while true
        push!(funcstrs, replace(strs[k], "#" => ""))
        k += 1
        while isnothing(findfirst(';', funcstrs[end]))
            funcstrs[end] = funcstrs[end] * strs[k]
            k += 1
        end
        push!(Ts, strs[k])
        k += 1
        if strs[k] == "N"
            k += 1
            break
        elseif strs[k] == "Y"
            k += 1
            continue
        else
            error("unexpected format in the tdb file")
        end
    end

    N = length(funcstrs)
    str = "function $(name)(T)::typeof(T)\n"
    for i in 1:N
        funcstrs[i] = strip(funcstrs[i], ';')

        str *= i == 1 ? "\tif " : "\telseif "
        str *= "$(Ts[i]) ≤ T ≤ $(Ts[i+1])\n"
        str *= "\t\t" * funcstrs[i] * "\n"
    end
    str *= "\tend\nend\n"

    Tl, Tu = tryparse.(Int, [Ts[1], Ts[end]])
    if nothing in [Tl, Tu]
        Tl, Tu = parse.(Float64, [Ts[1], Ts[end]])
    end

    return (Tl, Tu), str
end

function parse_typedefinition(text::AbstractString)
    strs = split(text)
    @assert match(r"TYPE_DEFINITION"i, strs[1]) ≠ ""
    code = strs[2]
    @assert length(code) == 1
    code = code[1]
    key = strs[3]
    if !(key in ["SEQ", "GES"])
        @warn "Unsupported keyword, $key, for TYPE_DEFINITION"
    end
    args = join(strs[4:end], ' ')
    return TypeDefinition(code, key, args)
end

function parse_phase(text::AbstractString)
    strs = split(text)
    @assert match(r"PHASE"i, strs[1]) ≠ ""
    model = strs[3] # "%", "%&", "%Q+", etc.
    K = parse(Int, strs[4])
    i = 5 + K - 1
    sites = tryparse.(Int, strs[5:i])
    if nothing in sites
        sites = parse.(Float64, strs[5:i])
    end
    text = replace(strs[2], " " => "")
    strs = split(text, ':')
    @assert length(strs) ≤ 2
    name = strs[1]
    state = length(strs) == 1 ? 'S' : strs[2][1]
    return Phase(name, state[1], model, sites)
end

function parse_constituent(text::AbstractString)
    strs = split(text)
    @assert match(r"CONSTITUENT"i, strs[1]) ≠ ""
    phas = split(strs[2], ':')[1] # "Liquid", "Bcc", "Fcc", etc.
    cons_text = join(strs[3:end]) # ":Cu,Zn", ":Cu,Zn:Cu,Zn:", etc.
    cons_text = format_constitution(cons_text)
    cons_strs = split(cons_text, ':')
    cons_strs = filter(e -> e ≠ "", cons_strs)
    cons = map(text -> split(text, ','), cons_strs)
    return phas, cons
end

function parse_parameter(text::AbstractString)
    # Parameters
    strs = split(text)
    @assert match(r"PAR"i, strs[1]) ≠ ""

    text = join(strs[2:end], ' ') # ex. G(BCC2,...

    k = findfirst('(', text)
    l = findfirst(')', text)
    name = replace(text[1:l], " " => "") # G(BCC_B2,Cu:Cu,Zn;0)
    name = format_constitution(name)
    funcname = getfuncname(name) # GBCC_B2CuCuZn0
    symbol = split(name, '(')[1] # G

    # Skip MatCalc specific parameters
    symbol == "HMVA" && return nothing, nothing
    symbol == "SE" && return nothing, nothing

    symbol = symbol[1]

    strs = split(text[l+1:end]) # 298.15...

    text = text[k+1:l-1] # BCC_B2,Cu,Zn:Cu,Zn;0
    text = format_constitution(text)

    k = findfirst(',', text)
    phas = replace(text[1:k-1], " " => "") # BCC_B2

    text = text[k+1:end] # Cu:Cu,Zn;0
    name = symbol * '(' * text * ')'
    comb_text, order_text = split(text, ';') # Cu:Cu,Zn, 0
    comb = parse_combination(comb_text)
    order = parse(Int, order_text)

    temp, funcstr = parse_function(funcname, strs)

    return phas, Parameter(name, symbol, comb, order, temp, funcname, funcstr)
end

function parse_combination(str::AbstractString)
    comb_strs = split(str, ':') # ["Cu", "Cu,Zn"]
    comb = map(text -> split(text, ','), comb_strs) # [[Cu], [Cu,Zn]]
end

function getfuncname(str::AbstractString)
    filter(c -> !in(c, ['(', ')', ',', ':', ';']), str)
end

function localname(param::Parameter)
    str = ""
    S = length(param.comb)
    for s in 1:S
    latt = param.comb[s]
        for cons in latt
            str *= subscript(cons)
            if cons ≠ latt[end]
                str *= ','
            end
        end
        if s ≠ S
            str *= ':'
        end
    end
    str *= ';' * string(param.order)
end

function localname(func::GFunction)
    return func.name
end

function format_constitution(str::AbstractString)
    # ex. str = "G(FCC,Cu:Cu,Zn,P;0)"
    replace(str, r"(?<=,|;|:).+?(?=,|;|:)" => format_specie)
end

function format_specie(str::AbstractString)
    # ex. str = "Cu2S1", "CU2O"
    str = replace(str, "%" => "")
    str = replace(str, r"(?<=\D)1(?=\D|$)" => "")
    replace(str, r"\D{1,2}(?=\d{1,}|$)" => format_element)
end

function format_element(str::AbstractString)
    N = length(str)
    N > 2 && error("Invalid element name: $str")
    N == 2 && return uppercase(str[1]) * lowercase(str[2])
    N == 1 && return uppercase(str)
end


using Test
using Dorothy

@DorothyAll()

@testset "Topology" begin

	@testset "Inference" begin
		m1 = readf("$(Dorothy.datapath)/1BTL.gro")
		topo = infertopology!(m1)
		@test topo[1] == [2]
		@test topo[2] == [1,3,5]
		@test topo[3] == [2,4,11]
		@test topo[4] == [3]
		m2 = readf("$(Dorothy.datapath)/MHC.pdb")
		topo = infertopology!(m2)
		@test topo[1] == [2,3,4,5,6,7]
		@test topo[2] == [1,3,5,6,7]
		@test topo[3] == [1,2,4,5,6,7]
		@test topo[4] == [1,3]
		@test topo[6540] == [6541, 6542]
		@test length(fragments(view(m2, Protein))) == 3
	end

end # @testset


module Grids

export Center, Face
export AbstractTopology, Periodic, Bounded, Flat, Connected, topology

export AbstractGrid, AbstractUnderlyingGrid, halo_size, total_size
export AbstractRectilinearGrid, RectilinearGrid
export XRegRectilinearGrid, YRegRectilinearGrid, ZRegRectilinearGrid, HRegRectilinearGrid, RegRectilinearGrid
export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid
export LatitudeLongitudeGrid, XRegLatLonGrid, YRegLatLonGrid, ZRegLatLonGrid
export ConformalCubedSphereFaceGrid, ConformalCubedSphereGrid
export node, xnode, ynode, znode, xnodes, ynodes, znodes, nodes
export offset_data, new_data
export on_architecture

using CUDA
using Adapt
using OffsetArrays

using Oceananigans
using Oceananigans.Architectures

import Base: size, length, eltype, show
import Oceananigans: short_show
import Oceananigans.Architectures: architecture

#####
##### Abstract types
#####

"""
    Center

A type describing the location at the center of a grid cell.
"""
struct Center end

"""
	Face

A type describing the location at the face of a grid cell.
"""
struct Face end

"""
    AbstractTopology

Abstract supertype for grid topologies.
"""
abstract type AbstractTopology end

"""
    Periodic

Grid topology for periodic dimensions.
"""
struct Periodic <: AbstractTopology end

"""
    Bounded

Grid topology for bounded dimensions, e.g., wall-bounded dimensions.
"""
struct Bounded <: AbstractTopology end

"""
    Flat

Grid topology for flat dimensions, generally with one grid point, along which the solution
is uniform and does not vary.
"""
struct Flat <: AbstractTopology end

"""
    Connected

Grid topology for dimensions that are connected to other models or domains on both sides.
"""
const Connected = Periodic  # Right now we just need them to behave like Periodic dimensions except we change the boundary conditions.

"""
    AbstractGrid{FT, TX, TY, TZ}

Abstract supertype for grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractUnderlyingGrid{FT, TX, TY, TZ}

Abstract supertype for "primary" grids (as opposed to grids with immersed boundaries)
with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} <: AbstractGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractRectilinearGrid{FT, TX, TY, TZ}

Abstract supertype for rectilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractRectilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractCurvilinearGrid{FT, TX, TY, TZ}

Abstract supertype for curvilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractUnderlyingGrid{FT, TX, TY, TZ, Arch} end

"""
    AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ}

Abstract supertype for horizontally-curvilinear grids with elements of type `FT` and topology `{TX, TY, TZ}`.
"""
abstract type AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch} <: AbstractCurvilinearGrid{FT, TX, TY, TZ, Arch} end

include("grid_utils.jl")
include("zeros.jl")
include("new_data.jl")
include("automatic_halo_sizing.jl")
include("input_validation.jl")
include("grid_generation.jl")
include("rectilinear_grid.jl")
include("conformal_cubed_sphere_face_grid.jl")
include("latitude_longitude_grid.jl")

end


__precompile__()

module Celeste

# import this before StaticArrays has a chance to load, to
# workaround a bug
import ForwardDiff

include("aliasscopes.jl")
include("config.jl")

# submodules
include("Log.jl")

include("SensitiveFloats.jl")
include("BivariateNormals.jl")

include("Coordinates.jl")

include("Model.jl")
include("Infer.jl")
include("Transform.jl")
include("PSF.jl")
include("SDSSIO.jl")
include("SEP.jl")

include("MCMC.jl")
include("StochasticVI.jl")
include("DeterministicVI.jl")
include("ParallelRun.jl")

include("Synthetic.jl")
include("AccuracyBenchmark.jl")
include("GalsimBenchmark.jl")

include("ArgumentParse.jl")

include("settings.jl")

import .ParallelRun: BoundingBox, OptimizedSource
import .SDSSIO: RunCamcolField, CatalogEntry
export BoundingBox, OptimizedSource, RunCamcolField, CatalogEntry

end # module


module Layer2Tests

using Test
using HDF5
using RandomMatrices
using LinearAlgebra

using PicoQuant

include("test_utils.jl")

@testset "Layer 2 tests" begin
    """
        function dsl_backend_teardown()

    Function to cleanup after the dslbackend
    """
    function dsl_backend_teardown()
        rm("contract_network.tl", force=true)
        rm("tensor_data.h5", force=true)
    end

    """Dictionary of init and finalise methods for testing backends"""
    backends = Dict("DSLBackend" =>
                    Dict(:init => () -> DSLBackend("contract_network.tl", "tensor_data.h5", "", true),
                        # :execute => (args...) -> execute_dsl_file(ComplexF64, args...),
                        :execute => execute_dsl_file,
                        :finalise => dsl_backend_teardown),

                    "InteractiveBackend" =>
                    Dict(:init => InteractiveBackend,
                        :execute => () -> (),
                        :finalise => () -> ())
                    )

    for (backend_name, backend_funcs) in pairs(backends)

        @testset "Test contract functions for $backend_name" begin

            # Create a tensor network and contraction plan to work with.
            qasm_str = """OPENQASM 2.0;
                        include "qelib1.inc";
                        qreg q[3];
                        h q[0];
                        cx q[0],q[1];
                        cx q[1],q[2];"""

            try
                circ = load_qasm_as_circuit(qasm_str)
                tng = convert_qiskit_circ_to_network(circ, backend_funcs[:init]())
                add_input!(tng, "000")
                add_output!(tng, "000")
                plan = random_contraction_plan(tng)

                # Test if contract_network creates files with dsl and tensor data.
                contract_network!(tng, plan)

                backend_funcs[:execute]()

                # Is there only one node left after the contraction?
                @test begin
                    length(tng.nodes) == 1
                end
            finally
                backend_funcs[:finalise]()
            end

            try
                # Create a network which will be disjoint when all edges are contracted
                qasm_str = """OPENQASM 2.0;
                            include "qelib1.inc";
                            qreg q[2];
                            h q[0];
                            h q[1];"""

                circ = load_qasm_as_circuit(qasm_str)
                tng = convert_qiskit_circ_to_network(circ, backend_funcs[:init]())
                add_input!(tng, "00")
                plan = random_contraction_plan(tng)

                contract_network!(tng, plan, "vector")

                # Is there only one node left in nodes array after the contraction?
                @test begin
                    length(tng.nodes) == 1
                end

                backend_funcs[:execute]()

                # Check if the final tensor is a vector
                @test begin
                    result = load_tensor_data(tng, :result)
                    length(size(result)) == 1
                end

                @test begin
                    result = load_tensor_data(tng, :result)
                    real(result)[1] ≈ 1/2
                end
            finally
                backend_funcs[:finalise]()
            end
        end

        @testset "Test full wavefunction contraction for $backend_name" begin

            # Create a tensor network to work with.
            qasm_str = """OPENQASM 2.0;
                        include "qelib1.inc";
                        qreg q[3];
                        h q[0];
                        cx q[0],q[1];
                        cx q[0],q[1];
                        cx q[0],q[2];
                        """

            try
                circ = load_qasm_as_circuit(qasm_str)
                tn = convert_qiskit_circ_to_network(circ, backend_funcs[:init]())
                add_input!(tn, "000")

                # See if full wavefunction contraction completes.
                full_wavefunction_contraction!(tn, "vector")

                # check if the computed metrics are correct
                @test begin
                    max_size = tn.backend.metrics.max_tensor_size == 8
                    total_space = tn.backend.metrics.total_space_allocated == 44
                    flops = tn.backend.metrics.flops == 124
                    max_size && total_space && flops
                end

                backend_funcs[:execute]()

                # Is there only one node left after the contraction?
                @test begin
                    length(tn.nodes) == 1
                end
            finally
                backend_funcs[:finalise]()
            end
        end

        @testset "Test netcon contraction function for $backend_name" begin
            # Create a tensor network to work with.
            qasm_str = """OPENQASM 2.0;
                        include "qelib1.inc";
                        qreg q[3];
                        h q[0];
                        cx q[0],q[1];
                        cx q[0],q[2];
                        """

            try
                circ = load_qasm_as_circuit(qasm_str)
                tn = convert_qiskit_circ_to_network(circ, backend_funcs[:init]())
                add_input!(tn, "000")
                add_output!(tn, "000")

                # See if full wavefunction contraction completes.
                netcon_contraction!(tn)

                backend_funcs[:execute]()

                # Is there only one node left after the contraction?
                @test begin
                    length(tn.nodes) == 1
                end
            finally
                backend_funcs[:finalise]()
            end
        end

        @testset "Test bgreedy contraction function for $backend_name" begin
            # Create a tensor network to work with.
            qasm_str = """OPENQASM 2.0;
                        include "qelib1.inc";
                        qreg q[3];
                        h q[0];
                        cx q[0],q[1];
                        cx q[0],q[2];
                        """

            try
                circ = load_qasm_as_circuit(qasm_str)
                tn = convert_qiskit_circ_to_network(circ, backend_funcs[:init]())
                add_input!(tn, "000")
                add_output!(tn, "000")

                # See if full wavefunction contraction completes.
                bgreedy_contraction!(tn)

                backend_funcs[:execute]()

                # Is there only one node left after the contraction?
                @test begin
                    length(tn.nodes) == 1
                end
            finally
                backend_funcs[:finalise]()
            end
        end

        @testset "Test full wavefunction and MPS contractions match for $backend_name" begin

            try
                circ = create_simple_preparation_circuit(3, 1).compose(create_ghz_preparation_circuit(3))

                # simulate with full wf approach
                tn = convert_qiskit_circ_to_network(circ, backend_funcs[:init](), decompose=true, transpile=true)
                add_input!(tn, "0"^tn.number_qubits)

                # See if full wavefunction contraction completes.
                full_wavefunction_contraction!(tn, "vector")

                backend_funcs[:execute]()

                full_wf = load_tensor_data(tn, :result)

                tn = convert_qiskit_circ_to_network(circ, backend_funcs[:init](), decompose=true, transpile=true)
                add_input!(tn, "0"^tn.number_qubits)

                # See if full wavefunction contraction completes.
                mps_nodes = contract_mps_tensor_network_circuit!(tn)
                calculate_mps_amplitudes!(tn, mps_nodes)

                backend_funcs[:execute]()

                mps_wf = load_tensor_data(tn, :result)

                # Is there only one node left after the contraction?
                @test begin
                    mps_wf ≈ full_wf
                end
            finally
                backend_funcs[:finalise]()
            end
        end

        @testset "Test tensor chain compression for $backend_name" begin

            @test begin
                try
                    tn = TensorNetworkCircuit(2, backend_funcs[:init]())
                    add_input!(tn, "00")
                    compress_tensor_chain!(tn, collect(keys(tn.nodes)))
                    length(tn.nodes) == 2
                finally
                    backend_funcs[:finalise]()
                end
            end

            @testset begin

                try
                    # create a network and apply a ghz state preparation
                    # circuit. Then apply two additional cnot gates which should
                    # cancel out. At this point there should be 3 virtual bonds
                    # each with dimension 2. After a compression operation, this
                    # will be reduced to a single bond with dimension 2
                    qasm_str = """OPENQASM 2.0;
                                include "qelib1.inc";
                                qreg q[2];
                                h q[0];
                                cx q[0],q[1];
                                cx q[0],q[1];
                                cx q[0],q[1];
                                """
                    circ = load_qasm_as_circuit(qasm_str)
                    tn = convert_qiskit_circ_to_network(circ, backend_funcs[:init]())
                    add_input!(tn, "0"^2)
                    decompose_tensor!(tn, :node_2, [:index_3, :index_4],
                                    [:index_2, :index_5])
                    decompose_tensor!(tn, :node_3, [:index_4, :index_6],
                                    [:index_5, :index_7])
                    decompose_tensor!(tn, :node_4, [:index_6, :index_8],
                                    [:index_7, :index_9])

                    # contract all links except the virtual and open edges
                    plan = [k for (k, v) in pairs(tn.edges) if !v.virtual &&
                                                            v.src !== nothing &&
                                                            v.dst !== nothing]

                    for edge in plan
                        contract_pair!(tn, edge)
                    end

                    compress_tensor_chain!(tn, [:node_18, :node_19])

                    idx = findfirst(x -> x == :index_14, tn.nodes[:node_18].indices)

                    if backend_name == "DSLBackend"
                        save_output(tn, :node_18, "node_18")
                    end

                    backend_funcs[:execute]()

                    result = load_tensor_data(tn, :node_18)

                    size(result)[idx] == 2
                finally
                    backend_funcs[:finalise]()
                end
            end
        end

        @testset "Test tensor network contraction with MPS, full vector for $backend_name" begin
            @test begin
                try
                    n = 5
                    circ = create_ghz_preparation_circuit(n)
                    tn = convert_qiskit_circ_to_network(circ,
                                                        backend_funcs[:init](),
                                                        decompose=true,
                                                        transpile=true)
                    add_input!(tn, "0"^n)
                    mps_nodes = contract_mps_tensor_network_circuit!(tn)
                    calculate_mps_amplitudes!(tn, mps_nodes)
                    backend_funcs[:execute]()
                    ref_output = zeros(ComplexF64, 2^n)
                    ref_output[[1, end]] .= 1/sqrt(2)
                    load_tensor_data(tn, :result) ≈ ref_output
                finally
                    backend_funcs[:finalise]()
                end
            end
        end

        @testset "Test tensor network contraction with MPS, select amplitudes for $backend_name" begin
            @test begin
                try
                    n = 5
                    circ = create_ghz_preparation_circuit(n)
                    tn = convert_qiskit_circ_to_network(circ,
                                                        backend_funcs[:init](),
                                                        decompose=true,
                                                        transpile=true)
                    add_input!(tn, "0"^n)
                    mps_nodes = contract_mps_tensor_network_circuit!(tn)
                    backend_funcs[:execute]()
                    mps_state = MPSState(tn, mps_nodes)
                    size(mps_state) == Tuple(ones(4) .* 2)
                    length(mps_state) == 2 << n
                    mps_state["11111"] ≈ mps_state["00000"] ≈ 1/sqrt(2)
                    mps_state["10101"] ≈ 0
                finally
                    backend_funcs[:finalise]()
                end
            end
        end

        @testset "Test decopmose with threshold option for backend $backend_name" begin
            @test begin
                try
                    # create distribution to sample from to get a random unitary matrix
                    dist = Haar(1)
                    d = 2
                    random_gate_data = reshape(rand(dist, d^2), (d, d, d, d))

                    tn = TensorNetworkCircuit(2, backend_funcs[:init]())
                    add_gate!(tn, random_gate_data, [1, 2])

                    threshold = 0.5
                    # decompose gate between qubits with threshold of 0.5
                    new_nodes = decompose_tensor!(tn,
                                                :node_1,
                                                [:index_1, :index_3],
                                                [:index_2, :index_4],
                                                threshold=threshold)
                    save_output(tn, new_nodes[1], String(new_nodes[1]))

                    backend_funcs[:execute]()

                    # decompose gate data manually and check threshold was applied
                    F = svd(reshape(permutedims(random_gate_data, (1, 3, 2, 4)), (d^2, d^2)))

                    S_norm = sqrt(sum(F.S .^ 2))
                    expected_dim = sum(F.S ./ S_norm .> threshold)
                    virtual_bonds = virtualedges(tn, new_nodes[1])
                    bond_idx = findfirst(x -> x == virtual_bonds[1], tn.nodes[new_nodes[1]].indices)
                    size(load_tensor_data(tn, new_nodes[1]))[bond_idx] == expected_dim
                finally
                    backend_funcs[:finalise]()
                end
            end

            @testset "Test decopmose with max_rank option for backend $backend_name" begin
                @test begin
                    try
                        # create distribution to sample from to get a random unitary matrix
                        dist = Haar(1)
                        d = 2
                        random_gate_data = reshape(rand(dist, d^2), (d, d, d, d))

                        tn = TensorNetworkCircuit(2, backend_funcs[:init]())
                        add_gate!(tn, random_gate_data, [1, 2])

                        max_rank = 1
                        # decompose gate between qubits with threshold of 1.0
                        new_nodes = decompose_tensor!(tn,
                                                    :node_1,
                                                    [:index_1, :index_3],
                                                    [:index_2, :index_4],
                                                    max_rank=max_rank)
                        save_output(tn, new_nodes[1], String(new_nodes[1]))

                        backend_funcs[:execute]()

                        virtual_bonds = virtualedges(tn, new_nodes[1])
                        bond_idx = findfirst(x -> x == virtual_bonds[1], tn.nodes[new_nodes[1]].indices)
                        size(load_tensor_data(tn, new_nodes[1]))[bond_idx] == max_rank
                    finally
                        backend_funcs[:finalise]()
                    end
                end
            end

            @testset "Test tensor slicing backend $backend_name" begin
                @test begin
                    try
                        n = 4
                        circ = create_simple_preparation_circuit(n, 2).compose(create_qft_circuit(n))

                        # contract the regular way
                        tn = convert_qiskit_circ_to_network(circ,
                                                            backend_funcs[:init](),
                                                            decompose=true,
                                                            transpile=true)
                        add_input!(tn, "0"^n)
                        full_wavefunction_contraction!(tn, "vector")
                        backend_funcs[:execute]()
                        wf = load_tensor_data(tn, :result)

                        # now decompose to 4 partitions and contract each partition separately
                        N = 4
                        wfs = []
                        for p in 1:N
                            tn = convert_qiskit_circ_to_network(circ,
                                                                backend_funcs[:init](),
                                                                decompose=true,
                                                                transpile=true)
                            add_input!(tn, "0"^n)
                            bond_labels, bond_values = partition_network_on_virtual_bonds(tn, N, p)
                            slice_tensor_network(tn, bond_labels, bond_values)
                            full_wavefunction_contraction!(tn, "vector")
                            backend_funcs[:execute]()
                            push!(wfs, load_tensor_data(tn, :result))
                        end
                        sum(wfs) ≈ wf
                    finally
                        backend_funcs[:finalise]()
                    end
                end
            end
        end
    end
end
end