import pandas as pd
import numpy as np
from agcounts.extract import get_counts

def get_counts_csv(file, freq: int, epoch: int, fast: bool = True, verbose: bool = False, time_column: str = None):
  if verbose:
    print("Reading in CSV", flush = True)
  raw = pd.read_csv(file, skiprows=0)
  if time_column is not None:
    ts = raw[time_column]
    ts = pd.to_datetime(ts)
    time_freq = str(epoch) + "S"
    ts = ts.dt.round(time_freq)
    ts = ts.unique()
    ts = pd.DataFrame(ts, columns=[time_column])
  raw = raw[["X", "Y", "Z"]]
  if verbose:
    print("Converting to array", flush = True)  
  raw = np.array(raw)
  if verbose:
    print("Getting Counts", flush = True)    
  counts = get_counts(raw, freq = freq, epoch = epoch, fast = fast, verbose = verbose)
  del raw
  counts = pd.DataFrame(counts, columns = ["Axis1","Axis2","Axis3"])
  counts["AC"] = (counts["Axis1"]^2 + counts["Axis2"]^2 + counts["Axis3"]^2) ** 0.5
  ts = ts[0:counts.shape[0]]
  if time_column is not None:
    counts = pd.concat([ts, counts], axis=1)
  return counts

def convert_counts_csv(file, outfile, freq: int, epoch: int, fast: bool = True, verbose: bool = False, time_column: str = None):
  counts = get_counts_csv(file, freq = freq, epoch = epoch, verbose = True, time_column = time_column)
  counts.to_csv(outfile,index=False)
  return counts

counts = get_counts_csv("data.csv.gz", freq = 80, epoch = 60)
counts = convert_counts_csv("data.csv.gz", outfile="counts.csv.gz", 
freq=80, epoch = 60, verbose = True, time_column = "HEADER_TIMESTAMP")