Price Projections in Rust

main.rs

extern crate csv;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate peroxide;

use chrono::Local;
use clap::{App, Arg};
use csv::Reader;
use peroxide::fuga::*;
use rand::Rng;
use serde::Deserialize;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Write};

#[derive(Debug, Clone, Deserialize)]
pub struct Return {
  pub idx: i32,
  pub dt: String,
  pub adj_close: f64,
  pub ret: f64,
  pub next_ret: f64,
  pub abs_ret: f64,
  pub std_dev: f64,
  pub tag: String,
  pub tag_pattern: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ProjectedReturn {
  pub generation: i32,
  pub day: i32,
  pub tag_pattern: String,
  pub ret: f64,
  pub price: f64,
  pub timestamp: String,
  pub ticker: String,
  pub expiry_date: String,
}

fn main() {
  let matches = App::new("clotho")
    .version("v0.000")
    .author("choiway <waynechoi@gmail.com>")
    .about("Projected return generator")
    .arg(
      Arg::with_name("ticker_symbol")
        .short("t")
        .value_name("TICKER_SYMBOL")
        .help("Ticker symbol for security")
        .required(true),
    )
    .arg(
      Arg::with_name("filepath")
        .short("f")
        .value_name("FILEPATH")
        .help("Path to returns data for ticker")
        .required(true),
    )
    .arg(
      Arg::with_name("expiry_date")
        .short("x")
        .value_name("EXPIRY_DATE")
        .help("Expiry date of analysis")
        .required(true),
    )
    .arg(
      Arg::with_name("days_ahead")
        .short("d")
        .value_name("DAYS_AHEAD")
        .help("Number of days to project forward")
        .required(true),
    )
    .arg(
      Arg::with_name("generations")
        .short("g")
        .value_name("GENERATIONS")
        .help("The number of times to run projections")
        .required(true),
    )
    .get_matches();
  env_logger::init();
  let ticker_symbol: String = matches.value_of("ticker_symbol").unwrap().to_string();
  let filepath: &str = matches.value_of("filepath").unwrap();
  info!("[{}] from {}", ticker_symbol, filepath);
  // let ticker_symbol = "AAPL".to_string();
  let expiry_date: String = matches.value_of("expiry_date").unwrap().to_string();
  let generations: i32 = matches.value_of("generations").unwrap().parse().unwrap();
  let number_of_days: i32 = matches.value_of("days_ahead").unwrap().parse().unwrap();
  let std_window = 20;
  let returns = open_returns_csv(filepath);
  let returns_by_tag = init_returns_by_tag(&returns);
  let returns_by_tag_pattern = init_returns_by_tag_pattern(&returns);
  let last_n_returns = pop_tail_returns(&returns, std_window);
  // Should refactor naming here
  let init_window_returns: Vec<f64> = last_n_returns.into_iter().map(|x| x.ret).collect();
  let mut projected_returns: Vec<ProjectedReturn> = Vec::new();
  let local_time = Local::now().to_rfc3339();

  for gen in 0..generations {
    generate_projected_returns(
      gen,
      number_of_days,
      &mut projected_returns,
      returns[returns.len() - 1].clone(),
      init_window_returns.clone(),
      &returns_by_tag,
      &returns_by_tag_pattern,
      &ticker_symbol,
      &expiry_date,
      &local_time,
    );
    // print!(".");
    // io::stdout().flush().unwrap();

  }
  // println!("{:?}", projected_returns);
  info!("[{}] Writing csv", ticker_symbol);
  let filename = write_projected_returns(&projected_returns, ticker_symbol);
  print!("{}", filename);
}

fn generate_projected_returns(
  gen: i32,
  number_of_days: i32,
  projected_returns: &mut Vec<ProjectedReturn>,
  last_return: Return,
  window_returns: Vec<f64>,
  returns_by_tag: &HashMap<String, Vec<f64>>,
  returns_by_tag_pattern: &HashMap<String, Vec<f64>>,
  symbol: &String,
  expiry_date: &String,
  ts: &String,
) {
  let mut price = last_return.adj_close.clone();
  let mut tag = last_return.tag.clone();
  let mut tag_pattern = last_return.tag_pattern.clone();
  let mut window_returns = window_returns;
  let mut std = window_returns.sd();
  for day in 0..number_of_days {
    let current_ret =
      get_next_ret_by_tag_pattern(&returns_by_tag, &returns_by_tag_pattern, &tag_pattern, &tag);
    price = price * (1.0 + current_ret);
    tag = tag_today_return(current_ret, std);
    tag_pattern = update_tag_pattern(tag_pattern, tag.clone());
    let new_projected_return = ProjectedReturn {
      generation: gen + 1,
      day: day + 1,
      tag_pattern: tag.clone(),
      ret: current_ret,
      price: price,
      timestamp: ts.clone(),
      ticker: symbol.clone(),
      expiry_date: expiry_date.clone(),
    };
    // println!("{:?}", new_projected_return);
    projected_returns.push(new_projected_return);
    update_sample(&mut window_returns, current_ret);
    std = window_returns.sd();
  }
}

pub fn open_returns_csv(ticker_path: &str) -> Vec<Return> {
  let mut returns: Vec<Return> = vec![];
  let mut rdr = Reader::from_path(ticker_path).expect(&*format!("No file for {}", ticker_path));
  for row in rdr.records() {
    let record = row.unwrap();
    let r = Return {
      idx: record.get(0).unwrap().parse().unwrap(),
      dt: record.get(1).unwrap().to_string(),
      adj_close: record.get(2).unwrap().parse().unwrap(),
      ret: record.get(3).unwrap().parse().unwrap(),
      next_ret: record.get(4).unwrap().parse().unwrap(),
      abs_ret: record.get(5).unwrap().parse().unwrap(),
      std_dev: record.get(6).unwrap().parse().unwrap(),
      tag: record.get(7).unwrap().to_string(),
      tag_pattern: record.get(8).unwrap().to_string(),
    };
    returns.push(r);
  }
  returns
}

fn pop_tail_returns(returns: &Vec<Return>, len: usize) -> Vec<Return> {
  let final_length = returns.len().saturating_sub(len);
  let tail = returns.clone().split_off(final_length);
  tail
}

// Array os returns by the 3 day tag pattern
fn init_returns_by_tag_pattern(returns: &Vec<Return>) -> HashMap<String, Vec<f64>> {
  let mut returns_by_tag_pattern = HashMap::new();
  for ret in returns {
    returns_by_tag_pattern
      .entry(ret.tag_pattern.clone())
      .or_insert(vec![ret.next_ret])
      .push(ret.next_ret);
  }
  returns_by_tag_pattern
}

fn init_returns_by_tag(returns: &Vec<Return>) -> HashMap<String, Vec<f64>> {
  let mut returns_by_tag = HashMap::new();
  for ret in returns {
    returns_by_tag
      .entry(ret.tag.clone())
      .or_insert(vec![ret.next_ret])
      .push(ret.next_ret)
  }
  returns_by_tag
}

fn update_sample(sample: &mut Vec<f64>, new_ret: f64) {
  sample.remove(0);
  sample.push(new_ret);
}

fn update_tag_pattern(tag_pattern: String, tag: String) -> String {
  let mut pattern = tag_pattern;
  pattern.remove(0);
  pattern.push_str(&tag);
  pattern
}

fn tag_today_return(ret: f64, std: f64) -> String {
  let normal_std = ret / std;
  match normal_std {
    s if s >= 2.0 => "E".to_string(),
    s if s < 2.0 && s >= 1.0 => "C".to_string(),
    s if s < 1.0 && s >= 0.0 => "A".to_string(),
    s if s < 0.0 && s >= -1.0 => "B".to_string(),
    s if s < -1.0 && s >= -2.0 => "D".to_string(),
    s if s < -2.0 => "F".to_string(),
    _ => "A".to_string(),
  }
}

fn get_next_ret_by_tag_pattern(
  returns_by_tag: &HashMap<String, Vec<f64>>,
  returns_by_tag_pattern: &HashMap<String, Vec<f64>>,
  tag_pattern: &String,
  tag: &String,
) -> f64 {
  match returns_by_tag_pattern.get(tag_pattern) {
    Some(rets) => {
      let mut rng = rand::thread_rng();
      let index = rng.gen_range(0, rets.len());
      rets[index]
    }
    None => get_rand_next_return_by_tag(&returns_by_tag, &tag),
  }
}

fn get_rand_next_return_by_tag(returns_by_tag: &HashMap<String, Vec<f64>>, tag: &String) -> f64 {
  match returns_by_tag.get(tag) {
    Some(rets) => {
      let mut rng = rand::thread_rng();
      let index = rng.gen_range(0, rets.len());
      rets[index]
    }
    None => 0.0,
  }
}

pub fn write_projected_returns(projected_returns: &Vec<ProjectedReturn>, symbol: String) -> String {
  // debug!("writing projected returns to disk");
  let filename = format!("/tmp/{}_projected_returns.csv", symbol);
  let mut f = File::create(filename.clone()).expect("Unable to create file");
  for projected_return in projected_returns {
    let p = projected_return;
    write!(
      f,
      "{},{},{},{},{},{},{},{}\n",
      p.generation, p.day, p.tag_pattern, p.ret, p.price, p.timestamp, p.ticker, p.expiry_date
    )
    .unwrap();
  }
  filename
}