#!/usr/local/bin/perl
#
#Script to show projects in tree format
use strict;
use warnings;

use lib "/cell_root/software/HPCC/lib/site_perl";
#Temporary for debugging
#use lib "/afs/.glue.umd.edu/software/HPCC/1.1/lib/site_perl";

use Glue::HPCC::Cluster;
use Glue::HPCC::Project::Project;
use Glue::HPCC::Project::ProjectType;
use Glue::HPCC::Project::Department;
use Glue::HPCC::Project::College;

use Slurm::Sshare;

use Getopt::Long;

#Globals
my $inactive;

my %usage_by_project = ();

sub usage()
{	print STDERR <<EOF;
$0: Show projects in tree format

Usage:
$0 [ -verbose ] [ -cluster CLUSTER ] [ -root PARENT ] \
	[ -child CHILD ] [ --inactive ]

List projects in a "tree" format, showing parent-child
relationships.  

If CLUSTER is not provided, it will be defaulted based on the 
system you are running the command on (i.e. whatever HPC 
cluster it belongs to).

If PARENT is given, will show the PARENT project and all
of its children.  If CHILD is given, will show the projects
from which CHILD is descended.  Either or both may be given
and/or repeated, resulting in multiple trees.

The special value ROOT can be given as a parent, in which
case trees starting from all projects without a parent will be 
shown.

Normally, only active projects shown; with --inactive 
inactive projects will also be shown (with a * after the name).

Normally, only a list of project names is printed.  
If -verbose is given once, also gives sizes of allocations.
If -verbose is given twice, also gives current Slurm usage of allocations.

EOF
}

sub get_usage_for_project($$)
#Given hpc instance, and project name, gets usage.
#If project has a high priority allocation account, combines usage
#for both.  On success, returns usage in kSU.
{	my $hpc = shift;
	my $pname = shift;

	my $cluster = $hpc->db_clustername;

	my $proj = Glue::HPCC::Project::Project->project_by_name($hpc,
		$pname);
	die "Unable to lookup project $pname: $proj at "
		unless $proj && ref($proj);
	my $ptype = Glue::HPCC::Project::ProjectType->projecttype_by_pkey(
		$hpc, $proj->projtype);
	die "Unable to lookup ptype for proj $pname (type=" .
		$proj->projtype . "): $ptype at "
		unless $ptype && ref($ptype);

	my ($cpusec_used, $cpumin_limit, $cpumin_unused, $userdata, $rec);
	my $acct;
	my $used_cpusec = 0;
	return $used_cpusec unless $ptype->is_metered;

	if ( $ptype->has_normal_priority )
	{ 	$acct = $pname;
		$rec = Slurm::Sshare->usage_for_account_in_cluster(
			cluster=>$cluster, account=>$acct);
		unless ( $rec && ref($rec) eq 'ARRAY' )
		{	die "Error looking up usage for acct $acct: $rec at ";
		}
		($cpusec_used, $cpumin_limit, $cpumin_unused, $userdata) 
			= @$rec;
		$used_cpusec += $cpusec_used;
	}

	if ( $ptype->has_high_priority )
	{ 	$acct = "${pname}-hi";
		$rec = Slurm::Sshare->usage_for_account_in_cluster(
			cluster=>$cluster, account=>$acct);
		unless ( $rec && ref($rec) eq 'ARRAY' )
		{	die "Error looking up usage for acct $acct: $rec at ";
		}
		($cpusec_used, $cpumin_limit, $cpumin_unused, $userdata) 
			= @$rec;
		$used_cpusec += $cpusec_used;
	}


	my $used_kSU = $used_cpusec/60/60/1000;
	return $used_kSU;
}

sub get_all_roots($)
#Returns names of all projects w/out parents
{	my $hpc = shift;

	my $projs = Glue::HPCC::Project::Project->pdb_query(
		$hpc, parent=>undef);
	my @projs = @$projs;
	unless ( $inactive )
	{	@projs = grep { $_->active } @projs;
	}
	@projs = map { $_->name } @projs;
	return \@projs;
}

sub _tree_from_root($$;$);
sub _tree_from_root($$;$)
{	my $hpc = shift;
	my $root = shift;
	my $get_usage = shift;

	my $rootname = $root->name;
	my $rootid = $root->prj_pkey;
	my $children = Glue::HPCC::Project::Project->pdb_query(
		$hpc, parent => $rootid );
	$children = [] unless $children;
	die "Error looking up children of $rootname: $children at "
		unless ref($children);
	
	if ( $get_usage )
	{	my $rootusage = get_usage_for_project($hpc, $rootname);
		$usage_by_project{$rootname} = $rootusage;
	}

	my $branches;
	CHILD: foreach my $child (@$children)
	{	next CHILD unless $child->active || $inactive;
		my $branch = _tree_from_root($hpc, $child, $get_usage);
		$branches = [] unless $branches;
		push @$branches, $branch;
	}

	return [ $root, $branches ];
}

sub make_tree_from_root($$;$)
{	my $hpc = shift;
	my $projname = shift;
	my $get_usage = shift;

	my $proj = Glue::HPCC::Project::Project->project_by_name(
		$hpc, $projname);
	die "No project name $projname found in PDB.\n"
		unless $proj;
	die "Project $projname is inactive and --inactive not given.\n"
		unless ( $proj->active || $inactive );
	die "Error looking up $projname in PDB: $proj at "
		unless ref($proj);

	my $tree = _tree_from_root($hpc, $proj, $get_usage);
	return $tree;
} 

sub _print_tree_from_root($$$$);
sub _print_tree_from_root($$$$)
{	my $hpc = shift;
	my $tree = shift;
	my $verbose = shift;
	my $level = shift;

	my $prefix = "| " x $level;
	chop $prefix;
	$prefix .= '-' if $level;

	my $root = $tree->[0];
	my $branches = $tree->[1] || [];

	my $rootname = $root->name;
	$rootname .= '*' unless $root->active;
	if ( $verbose )
	{	my $ptid = $root->projtype;
		my $ptype = Glue::HPCC::Project::ProjectType->projecttype_by_pkey(
			$hpc, $ptid);
		my $ptname = $ptype->name;

		my $deptid = $root->department;
		my $dept = Glue::HPCC::Project::Department->dept_by_pkey(
			$hpc, $deptid);
		my $deptcode = $dept->code;

		my $collid = $dept->college;
		my $coll = Glue::HPCC::Project::College->college_by_pkey(
			$hpc, $collid);
		my $collcode = $coll->code;

		my $su = $root->su;
		my $ksu = $su/1000;
		my $ksu_used = $usage_by_project{$rootname} || 0;
		my $ksu_free = $ksu - $ksu_used;
		$ksu = sprintf("%0.3f kSU", $ksu);
		$ksu_used = sprintf("%0.3f kSU", $ksu_used);
		$ksu_free = sprintf("%0.3f kSU", $ksu_free);
		my $childsu = 0;
		my $child_ksuused = 0;
		foreach my $br (@$branches)
		{	my $child = $br->[0];
			$childsu += $child->su;
			if ( $verbose > 1 )
			{	my $cname = $child->name;
				my $tmpused = $usage_by_project{$cname} || 0;
				$child_ksuused += $tmpused;
			}
		}
		my $childksu = $childsu/1000;
		my $child_ksufree = $childksu - $child_ksuused;
		$childksu = sprintf("%0.3f kSU", $childksu);
		my $childstr='';
		$child_ksuused = sprintf("%0.3f kSU", $child_ksuused);
		$child_ksufree = sprintf("%0.3f kSU", $child_ksufree);

		if ( scalar(@$branches) )
		{	$childstr = " (children $childksu)";
			if ( $verbose > 1 )
			{	$childstr = " (children $childksu, used=$child_ksuused, free=$child_ksufree)";
			}
		}

		my $fmtstr = "%s%s Type=%s Dept/Coll=%s/%s Size: %s%s\n";
		my @data = ($prefix, $rootname, $ptname, $deptcode, 
			$collcode, $ksu, $childstr);
		if ( $verbose > 1 )
		{   $fmtstr = 
			"%s%s Type=%s Dept/Coll=%s/%s Size: %s Used %s Free %s%s\n";
		    @data = ( $prefix, $rootname, $ptname, $deptcode,
			$collcode, $ksu, $ksu_used, $ksu_free, $childstr);
		}
		printf $fmtstr, @data;
	} else
	{	print $prefix, $rootname, "\n";
	}
	foreach my $branch (@$branches)
	{	_print_tree_from_root($hpc, $branch, $verbose, $level+1);
	}
}

sub print_tree_from_root($$$)
{	my $hpc = shift;
	my $tree = shift;
	my $verbose = shift;

	my $level = 0;
	_print_tree_from_root($hpc, $tree, $verbose, $level);
}

sub _tree_from_leaf($$;$);
sub _tree_from_leaf($$;$)
{	my $hpc = shift;
	my $leaf = shift;
	my $get_usage = shift;

	#Should NOT need to check for active status here;
	#an inactive parent should not have active children.

	my $leafname = $leaf->name;

	if ( $get_usage )
	{	my $leafusage = get_usage_for_project($hpc, $leafname);
		$usage_by_project{$leafname} = $leafusage;
	}

	my $parentid = $leaf->parent;
	return [ $leaf ] unless defined $parentid;
	

	my $parent = Glue::HPCC::Project::Project->project_by_pkey(
		$hpc, $parentid);
	die "Error: parent of $leafname ($parentid) not found in PDB at "
		unless $parent;
	die "Error looking up parent of $leafname: $parent at "
		unless ref($parent);
	
	my $branch = _tree_from_leaf($hpc, $parent, $get_usage);
	return [ $leaf, @$branch ];
}

sub make_tree_from_leaf($$;$)
{	my $hpc = shift;
	my $projname = shift;
	my $get_usage = shift;

	my $proj = Glue::HPCC::Project::Project->project_by_name(
		$hpc, $projname);
	die "No project name $projname found in PDB.\n"
		unless $proj;
	die "Project $projname inactive and --inactive not given.\n"
		unless $proj->active || $inactive;
	die "Error looking up $projname in PDB: $proj at "
		unless ref($proj);

	my $tree = _tree_from_leaf($hpc, $proj, $get_usage);
	return $tree;
} 

sub print_tree_from_leaf($$$)
{	my $hpc = shift;
	my $tree = shift;
	my $verbose = shift;

	my @tree = reverse @$tree;
	my $prefix = '';

	foreach my $proj (@tree)
	{	my $name = $proj->name;
		$name .= '*' unless $proj->active;
		my $ptid = $proj->projtype;
		my $ptype = Glue::HPCC::Project::ProjectType->projecttype_by_pkey(
			$hpc, $ptid);
		my $ptname = $ptype->name;

		my $deptid = $proj->department;
		my $dept = Glue::HPCC::Project::Department->dept_by_pkey(
			$hpc, $deptid);
		my $deptcode = $dept->code;

		my $collid = $dept->college;
		my $coll = Glue::HPCC::Project::College->college_by_pkey(
			$hpc, $collid);
		my $collcode = $coll->code;

		my $su = $proj->su;
		my $ksu = $su/1000;
		my $ksu_used = $usage_by_project{$name} || 0;
		my $ksu_free = $ksu - $ksu_used;
		$ksu = sprintf("%0.3f kSU", $ksu);
		$ksu_used = sprintf("%0.3f kSU", $ksu_used);
		$ksu_free = sprintf("%0.3f kSU", $ksu_free);

		if ( $verbose > 1 )
		{	printf "%s%s Dept/Coll=%s/%s Size: %s Used: %s Free: %s\n", 
				$prefix, $name, $ptname, $deptcode, 
				$collcode, $ksu, $ksu_used, $ksu_free;
		} elsif ( $verbose )
		{	printf "%s%s Dept/Coll=%s/%s Size: %s\n", 
				$prefix, $name, $ptname, $deptcode, 
				$collcode, $ksu;
		} else
		{	print $prefix, $name, "\n";
		}

		if ( $prefix )
		{	$prefix = "  $prefix";
		} else
		{	$prefix = "|-";
		}
	}
}


#-------------------------------------------------------------------
#		Main
#-------------------------------------------------------------------

my ($help, $verbose, $cluster);
my $roots = [];
my $children = [];
$verbose = 0;

my $res = GetOptions(
	'help!' => \$help,
	'verbose+' => \$verbose,
	'cluster=s' => \$cluster,
	'inactive!' => \$inactive,

	'root=s' => $roots,
	'child=s' => $children,
);

unless ( $res )
{	usage;
	die "\nError processing options.\n";
}
if ( $help )
{	usage;
	exit 0;
}
if ( scalar(@ARGV) )
{	usage;
	my $tmp = join ", ", @ARGV;
	die "\nError: extraneous arguments [ $tmp ], aborting.\n";
}

#Default cluster if needed
my $hpc = Glue::HPCC::Cluster->new_from_clustername($cluster);
die "\nInvalid cluster $cluster at " if $cluster && ! $hpc && ! ref($hpc);
die "\nError instantiating cluster $hpc " unless $hpc && ref($hpc);
$cluster = $hpc->db_clustername;

#Get all roots if ROOT is given
my @tmp = grep /ROOT/, @$roots;
my $tmp;
my %temp;
if ( @tmp )
{	$tmp = get_all_roots($hpc);
	push @$roots, @$tmp;
	%temp = map { $_ => undef } @$roots;
	delete $temp{ROOT};
	@tmp = keys %temp;
	@tmp = sort @tmp;
	$roots = [ @tmp ];
}

unless ( scalar(@$roots) || scalar(@$children) )
{	usage;
	die "\nYou must provide at least one parent or child.\n";
}

my $proj;
my $tree;
my $get_usage = 0;
$get_usage = 1 if $verbose > 1;

foreach $proj (@$roots)
{	$tree = make_tree_from_root($hpc, $proj, $get_usage);

	print_tree_from_root($hpc, $tree,$verbose);
}

foreach $proj (@$children)
{	$tree = make_tree_from_leaf($hpc, $proj, $get_usage);

	print_tree_from_leaf($hpc, $tree,$verbose);
}
