include("../helper/helperSystem.jl")

import DataFrames
import GLM
import Plots
import CSV

struct AmrexCastro <: AbstractAmrex
    app::String
    outputPrefix::String
    inputs::Array{String}
    outputs::Dict{String,Array{String}}
    runlogFile::String

    AmrexCastro(outputPrefix::String, runlogFile::String) = new(
        "AmrexCastro",
        outputPrefix,
        [
            "max_step",
            "amr.check_int",
            "amr.plot_int",
            "amr.n_cell",
            "amr.max_level",
            "castro.cfl",
            "castro.max_grid_size",
        ],
        Dict{String,Array{String}}(
            "plots_size" => ["amr.nplot_files", "amr.ncells"],
            "checks_size" => [
                "amr.check_int",
                "amr.ncheck_files",
                "amr.ncells",
                "amr.max_level",
                "amr.regrid_int",
            ],
        ),
        runlogFile,
    )
end

# runs the linear models for each entry in outputs
function run_linear_models(extractor::AmrexCastro)

    X::DataFrames.DataFrame = _get_linear_model_X(extractor)

    for key in keys(extractor.outputs)
        if (key == "plots_size")
            #_run_linear_model_plots_size(extractor, X)
            _run_linear_models_plots_size_step(extractor, X)
        end
    end
end


function _get_input_file(extractor::AmrexCastro, outputDir::String)::String

    # hardcoded input...find out if this is always the case with Castro   
    inputFiles::Array{String} = helper_get_prefix_files("input", outputDir)

    if isempty(inputFiles)
        throw(NoSuchFieldException(
            "jexio AmrexCastro: Can't find Castro input files in directories:",
            outputDirs,
        ))
        # use length, instead of size as it returns tuples
    elseif length(inputFiles) != 1
        print("Number of files: ", size(inputFiles))
        throw(DimensionMismatch("jexio AmrexCastro: found more than one input file"))
    else
    end

    # print(inputFiles)
    return inputFiles[1]
end



function _get_linear_model_X(extractor::AmrexCastro)::DataFrames.DataFrame
    # get directories runs output location
    outputDirs = helper_get_prefix_directories(extractor.outputPrefix)
    println(outputDirs)

    X::DataFrames.DataFrame = DataFrames.DataFrame()
    isXInit::Bool = false

    for outputDir in outputDirs
        inputFile::String = _get_input_file(extractor, outputDir)
        # get relevant input parameters
        parameters = _input_parser(extractor, inputFile)
        independentVariables = _get_independent_variables(extractor, parameters)
        # get outputDir name
        independentVariables["caseID"] = helper_get_relative_path(outputDir)

        if isXInit == false
            X = DataFrames.DataFrame(independentVariables)
            isXInit = true
        else
            DataFrames.push!(X, independentVariables)
        end
    end

    println(X)
    return X
end


function _run_linear_models_plots_size_step(extractor::AmrexCastro, X::DataFrames.DataFrame) 

  # using Flux ML for 1st order fitting
  # input with total_cells * output_number
  Xd::Array{Int64,1} = []
  # output directory size
  Yd::Array{Float64,1} = []
  # output Level0 directory size
  Yd_Level0::Array{Float64,1} = []
  Yd_Level1::Array{Float64,1} = []
  Yd_Level2::Array{Float64,1} = []
    
  caseIDs::Array{String,1} = []
  caseDirs::Array{String,1} = []

  outputDirs = helper_get_prefix_directories(extractor.outputPrefix)
      
  for (index,outputDir) in enumerate(outputDirs)
    
    inputFile::String = _get_input_file(extractor, outputDir)
    parameters = _input_parser(extractor, inputFile)
    rootPlotName = get(parameters, "amr.plot_file", "")
    # rootCheckpointName = get(parameters, "amr.plot_file", "")

    # find all directories with rootPlotName and get its size
    # root of _plt directories 
    plotFileDirs = helper_get_prefix_directories(string(outputDir, "/", rootPlotName))

    plotsSize::Int64 = 0
    plotsSizeLevel0::Int64 = 0
    plotsSizeLevel1::Int64 = 0
    plotsSizeLevel2::Int64 = 0
    
    counter::Int32 = 1
    
    # extract ncells
    ncells = X[!,"amr.ncells"][index]
    
    # caseID
    caseID = helper_get_relative_path(outputDir)
    
    for plotFileDir in plotFileDirs
      # Xd
      push!( Xd, counter*ncells)
      
      # Yd
      sizeDir::Int64 = helper_get_directory_size(plotFileDir)
      plotsSize += sizeDir
      push!( Yd, plotsSize)
      
      # Yd_Level0
      sizeDirLevel0::Int64 = helper_get_directory_size(string(plotFileDir, "/Level_0"))
      plotsSizeLevel0 += sizeDirLevel0
      push!( Yd_Level0, plotsSizeLevel0)
      
      # Yd_Level1
      sizeDirLevel1::Int64 = helper_get_directory_size(string(plotFileDir, "/Level_1"))
      plotsSizeLevel1 += sizeDirLevel1
      push!( Yd_Level1, plotsSizeLevel1)
      
      # Yd_Level2
      sizeDirLevel2::Int64 = helper_get_directory_size(string(plotFileDir, "/Level_2"))
      plotsSizeLevel2 += sizeDirLevel2
      push!( Yd_Level2, plotsSizeLevel2)
      
      
      # caseID
      push!(caseIDs, caseID )
      
      # caseDir
      push!(caseDirs, helper_get_relative_path(plotFileDir) )
      
      counter = counter + 1
    end
    
  end

  data = [(Xd,Yd)]
  
  df::DataFrames.DataFrame = 
      DataFrames.DataFrame( plots_cells = Xd, 
                            plots_size = Yd, 
                            plots_size_L0 = Yd_Level0,
                            plots_size_L1 = Yd_Level1,
                            plots_size_L2 = Yd_Level2,
                            caseID = caseIDs, 
                            caseDir = caseDirs)
  println(df)
  
  CSV.write("plot_size.csv",  df, header=true)
  #display(Plots.scatter(Xd, Yd))
  #readline()
end