#!/usr/bin/env julia
#############################################################################
## 00_validation.jl
## Pre-flight validation checks for Stan data and initialization
## Prevents cryptic Stan errors by catching issues early
#############################################################################

using Statistics

"""
Validate Stan data dictionary before passing to Stan.
Catches common issues that cause Stan to crash with cryptic errors.
"""
function validate_stan_data(dat::Dict; verbose=true)
    verbose && println("\n" * "="^70)
    verbose && println("VALIDATING STAN DATA")
    verbose && println("="^70)

    errors = String[]
    warnings = String[]

    # Check for NaN/Inf in all numeric data
    for (key, value) in dat
        if isa(value, AbstractArray) && eltype(value) <: Number
            if any(isnan, value)
                push!(errors, "Data '$key' contains NaN values")
            end
            if any(isinf, value)
                push!(errors, "Data '$key' contains Inf values")
            end
        elseif isa(value, Number)
            if isnan(value)
                push!(errors, "Data '$key' is NaN")
            end
            if isinf(value)
                push!(errors, "Data '$key' is Inf")
            end
        end
    end

    # Validate expert data is in open interval (0, 1)
    if haskey(dat, "val_dim")
        val_dim = dat["val_dim"]
        if any(x -> x <= 0 || x >= 1, val_dim)
            n_boundary = count(x -> x <= 0 || x >= 1, val_dim)
            push!(errors, "Expert dimension data has $n_boundary values at boundaries (must be in (0,1))")
            if verbose
                println("  Dimension-specific expert data range: [$(minimum(val_dim)), $(maximum(val_dim))]")
            end
        end
    end

    if haskey(dat, "val_lr")
        val_lr = dat["val_lr"]
        if any(x -> x <= 0 || x >= 1, val_lr)
            n_boundary = count(x -> x <= 0 || x >= 1, val_lr)
            push!(errors, "Expert L-R data has $n_boundary values at boundaries (must be in (0,1))")
            if verbose
                println("  L-R expert data range: [$(minimum(val_lr)), $(maximum(val_lr))]")
            end
        end
    end

    # Validate manifesto data
    if haskey(dat, "positive") && haskey(dat, "sample")
        positive = dat["positive"]
        sample = dat["sample"]

        if any(positive .> sample)
            push!(errors, "Manifesto: positive counts exceed sample sizes")
        end

        if any(positive .< 0)
            push!(errors, "Manifesto: negative positive counts found")
        end

        if any(sample .< 0)
            push!(errors, "Manifesto: negative sample sizes found")
        end
    end

    # Validate indices are within bounds
    # V10: Check segment indices (ss_man) if present, otherwise party indices (jj_man)
    if haskey(dat, "S") && haskey(dat, "ss_man")
        S = dat["S"]
        ss_man = dat["ss_man"]
        if any(ss_man .< 1) || any(ss_man .> S)
            push!(errors, "Manifesto segment indices out of bounds [1, $S]")
        end
    elseif haskey(dat, "J") && haskey(dat, "jj_man")
        J = dat["J"]
        jj_man = dat["jj_man"]
        if any(jj_man .< 1) || any(jj_man .> J)
            push!(errors, "Manifesto party indices out of bounds [1, $J]")
        end
    end

    if haskey(dat, "R") && haskey(dat, "rr_man")
        R = dat["R"]
        rr_man = dat["rr_man"]
        if any(rr_man .< 1) || any(rr_man .> R)
            push!(errors, "Manifesto party-year indices out of bounds [1, $R]")
        end
    end

    # V4: Validate constituent arrays
    if haskey(dat, "const_rr_man") && haskey(dat, "R")
        R = dat["R"]
        const_rr = dat["const_rr_man"]
        if any(const_rr .< 1) || any(const_rr .> R)
            push!(errors, "const_rr_man out of bounds [1, $R]")
        end
        # Verify offsets are valid
        if haskey(dat, "const_offset_man") && haskey(dat, "n_const_man")
            offsets = dat["const_offset_man"]
            n_consts = dat["n_const_man"]
            total = dat["N_const_man_total"]
            for i in eachindex(offsets)
                if offsets[i] + n_consts[i] - 1 > total
                    push!(errors, "const_offset_man[$i] + n_const_man[$i] exceeds N_const_man_total")
                    break
                end
            end
        end
    end

    # Print summary
    if verbose
        println("\nDATA SUMMARY:")
        # V10: Show segments if present
        if haskey(dat, "S")
            println("  Segments (S): $(dat["S"])")
            println("  Parties with segments (J): $(get(dat, "J", "N/A"))")
        else
            println("  Parties (J): $(get(dat, "J", "N/A"))")
        end
        println("  Countries (P): $(get(dat, "P", "N/A"))")
        println("  Segment-years (R): $(get(dat, "R", "N/A"))")
        println("  Years (T_year): $(get(dat, "T_year", "N/A"))")
        println("  Manifesto obs: $(get(dat, "N_man", "N/A"))")
        println("  Expert dim obs: $(get(dat, "N_exp_dim", "N/A"))")
        println("  Expert L-R obs: $(get(dat, "N_exp_lr", "N/A"))")

        if haskey(dat, "mn_resp_log_man")
            println("\nPRIOR MEANS:")
            println("  Manifesto: $(round(dat["mn_resp_log_man"], digits=3))")
            println("  Expert dim: $(round(dat["mn_resp_log_exp_dim"], digits=3))")
            println("  Expert L-R: $(round(dat["mn_resp_log_exp_lr"], digits=3))")
        end
    end

    # Report results
    if !isempty(errors)
        println("\n❌ VALIDATION FAILED - $(length(errors)) ERROR(S):")
        for (i, err) in enumerate(errors)
            println("  $i. $err")
        end
        return false
    end

    if !isempty(warnings)
        println("\n⚠️  $(length(warnings)) WARNING(S):")
        for (i, warn) in enumerate(warnings)
            println("  $i. $warn")
        end
    end

    if verbose
        println("\n✓ DATA VALIDATION PASSED")
        println("="^70)
    end

    return true
end

"""
Validate initialization values before passing to Stan.
Checks for common issues that cause immediate Stan crashes.
"""
function validate_init_values(init_dict::Dict; verbose=true)
    verbose && println("\n" * "="^70)
    verbose && println("VALIDATING INITIALIZATION VALUES")
    verbose && println("="^70)

    errors = String[]
    warnings = String[]

    for (key, value) in init_dict
        # Check for NaN/Inf
        if isa(value, AbstractArray) && eltype(value) <: Number
            if any(isnan, value)
                push!(errors, "Init '$key' contains NaN")
            end
            if any(isinf, value)
                push!(errors, "Init '$key' contains Inf")
            end

            if verbose && length(value) > 0
                val_array = vec(value)
                println("  $key: range [$(round(minimum(val_array), digits=3)), $(round(maximum(val_array), digits=3))]")
            end
        elseif isa(value, Number)
            if isnan(value)
                push!(errors, "Init '$key' is NaN")
            end
            if isinf(value)
                push!(errors, "Init '$key' is Inf")
            end

            if verbose
                println("  $key: $(round(value, digits=3))")
            end
        end

        # Check positive constraints (common Stan constraints)
        # Exception: *_raw parameters are non-centered and can be any real
        if (contains(string(key), "sigma") || contains(string(key), "tau") || contains(string(key), "phi")) &&
           !endswith(string(key), "_raw")
            if isa(value, Number) && value <= 0
                push!(errors, "Init '$key' = $value violates constraint > 0")
            elseif isa(value, AbstractArray) && any(value .<= 0)
                push!(errors, "Init '$key' has values ≤ 0 (violates constraint > 0)")
            end
        end

        # Check Cholesky factors are valid
        if contains(string(key), "L_Omega")
            if isa(value, AbstractMatrix)
                # Check it's lower triangular with positive diagonal
                n = size(value, 1)
                if size(value, 2) != n
                    push!(errors, "Init '$key' is not square")
                end
                for i in 1:n
                    if value[i, i] <= 0
                        push!(errors, "Init '$key' has non-positive diagonal at position $i")
                    end
                    for j in (i+1):n
                        if abs(value[i, j]) > 1e-10
                            push!(warnings, "Init '$key' is not lower triangular")
                            break
                        end
                    end
                end
            end
        end

        # Check slope parameters for positive constraint (V2/V3 feature)
        if key == "Gamma_man_slope_raw"
            if isa(value, AbstractArray) && any(value .< 0)
                push!(errors, "Init 'Gamma_man_slope_raw' has negative values (must be ≥ 0)")
            end
        end
    end

    # Report results
    if !isempty(errors)
        println("\n❌ INIT VALIDATION FAILED - $(length(errors)) ERROR(S):")
        for (i, err) in enumerate(errors)
            println("  $i. $err")
        end
        return false
    end

    if !isempty(warnings)
        println("\n⚠️  $(length(warnings)) WARNING(S):")
        for (i, warn) in enumerate(warnings)
            println("  $i. $warn")
        end
    end

    if verbose
        println("\n✓ INIT VALIDATION PASSED")
        println("="^70)
    end

    return true
end

"""
Estimate memory requirements for model
"""
function estimate_memory_requirements(dat::Dict; verbose=true, num_chains::Int=4, num_samples::Int=1000, num_threads_per_chain::Int=1)
    if !verbose
        return
    end

    println("\n" * "="^70)
    println("MEMORY ESTIMATE")
    println("="^70)

    R = get(dat, "R", 0)
    J = get(dat, "J", 0)
    K_man = get(dat, "K_man", 0)
    K_exp_dim = get(dat, "K_exp_dim", 0)
    K_exp_lr = get(dat, "K_exp_lr", 0)
    N_man = get(dat, "N_man", 0)
    N_ciy = get(dat, "N_ciy", 0)
    T_year = get(dat, "T_year", 0)

    # Rough parameter count
    theta_params = 4 * R
    item_params = 4 * K_man * 2 + K_exp_dim * 3 + K_exp_lr * 2
    strategic_params = get(dat, "P", 0) * get(dat, "K_man", 0)  # Country-item intercepts
    other_params = 4 * J + T_year + J + N_ciy + 50

    total_params = theta_params + item_params + strategic_params + other_params

    # Memory estimate (very rough)
    # Each parameter: ~8 bytes (float64) × samples × chains
    total_draws_per_param = num_samples * num_chains
    bytes_per_param = 8 * total_draws_per_param
    total_mb = (total_params * bytes_per_param) / (1024 * 1024)

    println("  Configuration:")
    println("    Chains: $num_chains")
    println("    Samples per chain: $num_samples")
    println("    Threads per chain: $num_threads_per_chain")
    println("    Total parallel workers: $(num_chains * num_threads_per_chain)")

    println("  Total parameters: ~$(total_params)")
    println("  Estimated memory (samples only): ~$(round(total_mb, digits=0)) MB")
    thread_scaling = max(1, num_threads_per_chain)
    println("  With thread overhead (×$(thread_scaling)): ~$(round(total_mb * thread_scaling, digits=0)) MB")
    println("  With safety margin (×3): ~$(round(3 * total_mb * thread_scaling, digits=0)) MB")

    if total_mb * thread_scaling * 3 > 8000
        println("\n⚠️  WARNING: Model may require > 8GB RAM")
    end

    println("="^70)
end
