From abda4bf7e5fb76111a2f8a0a1354d24a24c90bcd Mon Sep 17 00:00:00 2001
From: Jamie <jmaslek11@gmail.com>
Date: Sun, 11 Apr 2021 17:57:51 -0400
Subject: [PATCH] added select/mktcap options

---
 .../portfolio_optimization/port_opt_api.py    | 56 +++++++++----------
 .../port_optimization_controller.py           | 15 ++---
 2 files changed, 34 insertions(+), 37 deletions(-)

diff --git a/gamestonk_terminal/portfolio_optimization/port_opt_api.py b/gamestonk_terminal/portfolio_optimization/port_opt_api.py
index 862f45dc0bc..7c7e8d9f3d4 100644
--- a/gamestonk_terminal/portfolio_optimization/port_opt_api.py
+++ b/gamestonk_terminal/portfolio_optimization/port_opt_api.py
@@ -3,37 +3,9 @@ __docformat__ = "numpy"
 
 import argparse
 from typing import List
+import yfinance as yf
 
-def optimize(list_of_stocks:List[str], other_args:List[str]):
-    """
-
-    Parameters
-    ----------
-    list_of_stocks: List[str]
-        List of tickers to be included in optimization
-
-    Returns
-    -------
-    weights : dict
-        Dictionary of weights where keys are the tickers
-
-    """
-
-    parser = argparse.ArgumentParser(
-        add_help=False,
-        prog="optimize",
-    )
-
-    parse
-
-    weights = {}
-    n_stocks = len(list_of_stocks)
-    for stock in list_of_stocks:
-        weights[stock] = round(1/n_stocks, 5)
-
-    return weights
-
-def equal_weight(list_of_stocks:List[str]):
+def equal_weight(list_of_stocks:List[str], other_args:List[str]):
     """
         Equally weighted portfolio, where weight = 1/# of stocks
 
@@ -55,3 +27,27 @@ def equal_weight(list_of_stocks:List[str]):
 
     return weights
 
+def market_cap_weighting(list_of_stocks:List[str], other_args:List[str]):
+    """
+    Market cap weighted portfolio where each weight is the stocks market cap/ sum of all market caps
+    Parameters
+    ----------
+    list_of_stocks: List[str[
+        List of tickers to be included in optimization
+
+    Returns
+    -------
+    weights: dict
+        Dictionary of weights where keys are the tickers
+    """
+    weights = {}
+    mkt_cap = {}
+    total_cap = 0
+    for stock in list_of_stocks:
+        stock_cap = yf.Ticker(stock).info["marketCap"]
+        mkt_cap[stock] = stock_cap
+        total_cap += stock_cap
+    for k,v in mkt_cap.items():
+        weights[k] = round(v/total_cap,5)
+
+    return weights
diff --git a/gamestonk_terminal/portfolio_optimization/port_optimization_controller.py b/gamestonk_terminal/portfolio_optimization/port_optimization_controller.py
index 5892fb7f58c..aad015af3a5 100644
--- a/gamestonk_terminal/portfolio_optimization/port_optimization_controller.py
+++ b/gamestonk_terminal/portfolio_optimization/port_optimization_controller.py
@@ -86,11 +86,7 @@ class PortfolioOptimization:
     def call_add(self, other_args:List[str]):
         self.add_stocks(self, other_args)
 
-    def call_optimize(self, other_args:List[str]):
-        weights = po_api.optimize(self.tickers, other_args)
-        print("Optimal Weights for Equal Weighting:")
-        print(weights)
-        print("")
+
     def call_equal_weight(self, other_args:List[str]):
         weights = po_api.equal_weight(self.tickers, other_args)
         print("Optimal Weights for Equal Weighting:")
@@ -98,9 +94,14 @@ class PortfolioOptimization:
         print("")
 
     def call_mkt_cap(self,other_args:List[str]):
-        print("TODO")
+        weights = po_api.market_cap_weighting(self.tickers, other_args)
+        print("Market Cap Weighting Weights:")
+        print(weights)
+        print("")
+
     def call_select(self,other_args:List[str]):
-        print("TODO")
+        self.tickers = set([])
+        self.add_stocks(self, other_args)
 
     @staticmethod
     def add_stocks(self, other_args: List[str]):