#!/bin/bash

rankfile_generation() {
    
    # Check requested variables:
    if [[ -z "$NEMO_NODES" || -z "$XIOS_NODES" || -z "$RF_NAME" ]]; then
        echo "Error: Unable to find XIOS_NODES or NEMO_NODES or RF_NAME env variables!"
        return 1
    fi

    # Get job hosts from SLURM:
    JOB_HOSTS=$(scontrol show hostnames)
    local job_hosts=($JOB_HOSTS)

    local total_nodes=${#job_hosts[@]}
    local rank=0

    # Check total nodes
    if (( NEMO_NODES + XIOS_NODES > total_nodes )); then
        echo "Error: Less JOB_HOSTS for NEMO_NODES + XIOS_NODES."
        return 1
    fi

    # ENSURE RK EMPTY:
    > "$RF_NAME"
    # NEMO_NODES -> Hard-coded 112 MPI tasks per node:
    for (( i=0; i<NEMO_NODES; i++ )); do
        for (( j=0; j<112; j++ )); do
            local socket=$(( j / 56 ))  # Cambia de socket cada 56 cores
            local cpu=$(( j % 56 ))      # Usa los cores dentro del socket
            echo "rank $rank= ${job_hosts[$i]} slot=${socket}:$cpu" >> "$RF_NAME"
            ((rank++))
        done
    done

    # XIOS_NODES -> Hard-coded 4 MPI tasks per node:
    for (( i=NEMO_NODES; i<NEMO_NODES+XIOS_NODES; i++ )); do
        for (( j=0; j<4; j++ )); do
            echo "rank $rank= ${job_hosts[$i]} slot=0:0-27"  >> "$RF_NAME"; ((rank++))
            echo "rank $rank= ${job_hosts[$i]} slot=0:28-55" >> "$RF_NAME"; ((rank++))
            echo "rank $rank= ${job_hosts[$i]} slot=1:0-27"  >> "$RF_NAME"; ((rank++))
            echo "rank $rank= ${job_hosts[$i]} slot=1:28-55" >> "$RF_NAME"; ((rank++))

        done
    done

    echo "Succesfull generation of '$RF_NAME'."
}


rankfile_generation_logic() {
    
    # Check requested variables:
    if [[ -z "$NEMO_NODES" || -z "$XIOS_NODES" || -z "$RF_NAME" ]]; then
        echo "Error: Unable to find XIOS_NODES or NEMO_NODES or RF_NAME env variables!"
        return 1
    fi

    # Get job hosts from SLURM:
    JOB_HOSTS=$(scontrol show hostnames)
    local job_hosts=($JOB_HOSTS)

    local total_nodes=${#job_hosts[@]}
    local rank=0

    # Check total nodes
    if (( NEMO_NODES + XIOS_NODES > total_nodes )); then
        echo "Error: Less JOB_HOSTS for NEMO_NODES + XIOS_NODES."
        return 1
    fi

    # ENSURE RK EMPTY:
    > "$RF_NAME"
    # NEMO_NODES -> Hard-coded 112 MPI tasks per node:
    for (( i=0; i<NEMO_NODES; i++ )); do
        for (( j=0; j<112; j++ )); do
            echo "rank $rank= ${job_hosts[$i]} slot=$j" >> "$RF_NAME"
            ((rank++))
        done
    done

    # XIOS_NODES -> Hard-coded 4 MPI tasks per node:
    for (( i=NEMO_NODES; i<NEMO_NODES+XIOS_NODES; i++ )); do
        for (( j=0; j<XIOS_TPN; j++ )); do
            echo "rank $rank= ${job_hosts[$i]} slot=$j" >> "$RF_NAME"
            ((rank++))
        done
    done

    echo "Succesfull generation of '$RF_NAME'."
}

rankfile_generation_intel_mpi(){
    # Check requested variables:
    if [[ -z "$NEMO_NODES" || -z "$XIOS_NODES" || -z "$RF_NAME" ]]; then
        echo "Error: Unable to find XIOS_NODES or NEMO_NODES or RF_NAME env variables!"
        return 1
    fi
    # Get job hosts from SLURM:
    JOB_HOSTS=$(scontrol show hostnames)
    local job_hosts=($JOB_HOSTS)
    local total_nodes=${#job_hosts[@]}

    # Check total nodes
    if (( NEMO_NODES + XIOS_NODES > total_nodes )); then
        echo "Error: Less JOB_HOSTS for NEMO_NODES + XIOS_NODES."
        return 1
    fi

    # ENSURE RK EMPTY:
    > "$RF_NAME"

    # NEMO_NODES -> Hard-coded 112 MPI tasks per node:
    for (( i=0; i<NEMO_NODES; i++ )); do
        echo "${job_hosts[$i]}:112" >> $RF_NAME
    done

    # XIOS_NODES -> Hard-coded 4 MPI tasks per node:
    for (( i=NEMO_NODES; i<NEMO_NODES+XIOS_NODES; i++ )); do
        echo "${job_hosts[$i]}:$XIOS_TPN" >> $RF_NAME
    done

    echo "Succesfull generation of '$RF_NAME'."
}

