#!/usr/bin/env perl
#
# =============================================================================
# namsplit.pl: Splits a OASIS3 namcouple file for the pseudo-parallel mode
#
# UFla, 06/2010
# =============================================================================

use strict;
use warnings;

# Define regular expressions for some important lines of the namcouple file
my $re_fields_sec  = '^ \$STRINGS\s+$';
my $re_nfield_sec  = '^ \$NFIELDS\s+$';
my $re_field_begin = 'EXPORTED\s+$';
my $re_end_sec     = '^ \$END\s+$';
my $re_comment     = '^#';

# We need exactly two command line arguments:
# - the number of processors (i.e. the number of files to split into)
# - the namcouple file we want to split
@ARGV == 2 or die "Usage: namsplit.pl <NUM_PROCS> <NAMCOUPLE_FILE>\n";
(my $num_procs,my $namcouple) = @ARGV;

open NAMCOUPLE, "<$namcouple" or die "Can't open file '$namcouple'\n";

# The @header array holds all the header lines of the files, i.e. all
# lines _before_ any couple field definitions. Note that comments are
# stripped.
my @header = ();

# Now fill the @header by reading the namcouple file line-by-line
while (<NAMCOUPLE>)
{
    next if (/$re_comment/);           # Ignore comments
    push @header,$_;                   # Store current line in @header
    <NAMCOUPLE> if (/$re_nfield_sec/); # Ignore the line after the NFIELDS line
                                       #   (this line is recreated later for each
                                       #    split file)
    last if (/$re_fields_sec/);        # Stop if end of header found
}

# $num_fields counts the overall number of fields found in the namcouple file
my $num_fields = 0;

# The @fields array holds all the lines for each couple field as one string
my @fields;

# Count fields and fill @fields
while (<NAMCOUPLE>)
{
    next if (/$re_comment/); # Ignore comments
    last if (/$re_end_sec/); # Stop at $END token

    $num_fields++ if (/$re_field_begin/); # Increase the counter each
                                          #   time a new field is found

    $fields[$num_fields-1] .= $_; # Append the current line to the current
                                  #   field string in @fields
}

# Correct if we have more procs than fields
$num_procs = $num_fields if $num_fields<$num_procs;

# First distribution round: Assign equally many fields to each processor
my @fields_per_proc = map int($num_fields/$num_procs),(1..$num_procs);

# Second distribution round: Assign remaining number of fields one by one
foreach (1..$num_fields%$num_procs)
{
    $fields_per_proc[$_-1] += 1;
}

# Remember how many fields we have already written out
my $fld_counter = 0;

# Open and fill the split files one-by-one
for (my $i=0;$i<$num_procs;$i++)
{
    open OUT,"> $namcouple"."_".$i or die "Can't open output file\n";

    # Write the header first (line-by-line)
    foreach (@header)
    {
        print OUT $_;

        # Insert the correct number of coupling fields for each file
        print OUT "    $fields_per_proc[$i] $num_fields\n" if (/$re_nfield_sec/);
    }

    # Now write the coupling field definitions
    print OUT $fields[$fld_counter++] foreach (1..$fields_per_proc[$i]);

    # Write $END token and close the file...
    print OUT " \$END\n";
    close OUT;
}