#!/usr/bin/env python
import argparse
import pcsv.pindent
import pandas as pd
import sys
from jtutils import fix_broken_pipe, str_is_int, is_int, is_float, str_is_float

def readCL():
    parser = argparse.ArgumentParser()
    parser.add_argument("-f","--infile",type=argparse.FileType("r"),default=sys.stdin)
    parser.add_argument("-c","--sort_list",help="csv of column names or indices. Can include currently non-existent columns",default="")
    parser.add_argument("-r","--reverse",action="store_true",help="sort ascending (default descending)")
    parser.add_argument("-s","--string",action="store_true",help="sort as string (default is float)")
    parser.add_argument("--force",action="store_true",help="force sort, pushing invalid characters to the end")
    parser.add_argument("-d","--delimiter", default=",")
    parser.add_argument("-n","--no_header",action="store_true")
    parser.add_argument("-l","--lam", nargs="*", default=[], help="function body to sort by. For example: 'return str(x) or return (x-2)**2'")
    parser.add_argument("--sort_by_row",type=int,help="Sort the columns by a row instead of the other way around. Rows are zero-indexed.")
    args = parser.parse_args()

    if args.lam:
        for i,fn_string in enumerate(args.lam):
            fn_string = "def lam{i}(x): {fn_string}".format(**vars())
            exec(pcsv.pindent.pindent(fn_string))
        args.lam = vars()["lam0"]

    args.sort_list = args.sort_list.split(",")
    return args.infile, args.sort_list, args.reverse, args.string, args.delimiter, args.no_header, args.lam, args.sort_by_row, args.force


def sort_df(df, sort_list, string, reverse, lam, sort_by_row):
    if sort_by_row is not None:
        sort_list = [sort_by_row]
    else:
        sort_list = list(proc_sort_list(sort_list, df))

    if sort_by_row is not None:
        df = df.transpose() #transpose and then handle normally

    sort_fn = None
    if lam:
        sort_fn = lam
    elif not string:
        if force:
            last = float("inf") if reverse else float("-inf")
            sort_fn = lambda x: float(x) if str_is_float(x) else last
        else:
            def fn(x):
                try:
                    x = float(x)
                except:
                    raise Exception("Couldn't cast "+repr(x)+" to float. Use -s to string sort or --force option to ignore non-float values.")
                return x
            sort_fn = fn

    if string: #sort by string
        ascending = (not reverse)
    elif not lam: #sort by float
        ascending = reverse
    else: #custom sort function
        vals = (df[sort_list[0]].values[:5])
        if all(is_float(sort_fn(v)) for v in vals): #function returns float values
            ascending = reverse
        else:
            ascending = (not reverse)

    if sort_fn:
        tmp_df = pd.DataFrame()
        for s in sort_list:
            tmp_df[s] = df[s].apply(sort_fn)
        tmp_df = tmp_df.sort_values(by=sort_list, ascending=ascending)
        df = df.reindex(tmp_df.index)
    else:
        df = df.sort_values(by=sort_list, ascending=ascending)

    if sort_by_row is not None:
        df = df.transpose() #undo the transpose above
    return df


if __name__ == "__main__":
    infile, sort_list, reverse, string, delimiter, no_header, lam, sort_by_row, force = readCL()
    fix_broken_pipe()
    def proc_sort_list(sort_list, df):
        for index in sort_list:
            if index in df.columns:
                yield index
            elif str_is_int(index):
                yield df.columns[int(index)]
            else:
                raise Exception("ERROR: invalid sort_list element {index}".format(**vars()))
    args = {}
    if no_header:
        args["header"]=None
    df = pd.read_csv(infile, delimiter=delimiter, encoding="utf-8", dtype="object", **args)

    df = sort_df(df, sort_list, string, reverse, lam, sort_by_row)

    df.to_csv(sys.stdout,index=False,encoding="utf-8")
