Given a list of transactions, can we determine a strategy to match an actual cost basis?

cost_basis_puzzle.py

#!/usr/bin/env python3
from collections import namedtuple
from dataclasses import dataclass
import csv
import datetime
from typing import Sequence, Callable

__author__ = 'David Blume'
__copyright__ = "Copyright 2024, David Blume"
__license__ = "Apache 2.0"
__version__ = "0.1"


@dataclass
class Lot:
    date: str
    units: int
    unit_price: float


transaction_fields = ('type',
                      'trade_date',
                      'units',
                      'unit_price',
                      'total')


Transaction = namedtuple('Transaction', transaction_fields)

# A tuple of the actual transactions. These simulations
# will "replay" these.
transactions = (
    Transaction('buy', '2022-09-14', 601, 49.17, -29551.11),
    Transaction('buy', '2022-11-01', 4, 47.55, -190.19),
    Transaction('buy', '2022-11-07', 1, 47.6, -47.60),
    Transaction('buy', '2022-12-07', 1, 49.96, -49.96),
    Transaction('buy', '2022-12-09', 1, 49.95, -49.95),
    Transaction('buy', '2022-12-13', 3, 50.14, -150.41),
    Transaction('buy', '2023-01-17', 1, 50.54, -50.54),
    Transaction('buy', '2023-02-02', 3, 51.04, -153.14),
    Transaction('buy', '2023-02-07', 3, 50.6, -151.80),
    Transaction('buy', '2023-02-16', 1, 49.92, -49.93),
    Transaction('buy', '2023-03-07', 2, 49.44, -98.87),
    Transaction('buy', '2023-03-15', 431, 50.2, -21634.05),
    Transaction('buy', '2023-03-16', 3, 50.28, -150.83),
    Transaction('buy', '2023-03-17', 4, 50.3, -201.20),
    Transaction('buy', '2023-03-20', 290, 50.17, -14549.01),
    Transaction('buy', '2023-03-29', 5, 50.49, -252.44),
    Transaction('buy', '2023-03-30', 11, 50.49, -555.39),
    Transaction('buy', '2023-04-11', 2, 50.89, -101.77),
    Transaction('buy', '2023-04-17', 7, 50.72, -355.08),
    Transaction('buy', '2023-04-28', 206, 50.39, -10380.34),
    Transaction('buy', '2023-05-01', 3, 50.18, -150.54),
    Transaction('buy', '2023-05-19', 7, 49.78, -348.43),
    Transaction('buy', '2023-06-02', 5, 49.86, -249.30),
    Transaction('buy', '2023-06-05', 4, 49.9, -199.60),
    Transaction('buy', '2023-06-07', 1, 50.0, -50.00),
    Transaction('buy', '2023-06-12', 1, 49.93, -49.93),
    Transaction('buy', '2023-06-14', 10, 50.07, -500.70),
    Transaction('buy', '2023-06-16', 12, 50.16, -601.92),
    Transaction('buy', '2023-07-10', 1, 49.92, -49.93),
    Transaction('buy', '2023-07-11', 2, 49.97, -99.94),
    Transaction('buy', '2023-07-17', 1, 50.22, -50.22),
    Transaction('buy', '2023-07-26', 6, 50.42, -302.52),
    Transaction('buy', '2023-07-31', 5, 50.21, -251.03),
    Transaction('buy', '2023-08-01', 5, 49.95, -249.75),
    Transaction('buy', '2023-08-02', 5, 49.78, -248.88),
    Transaction('buy', '2023-08-03', 5, 49.48, -247.43),
    Transaction('buy', '2023-08-07', 7, 49.56, -346.92),
    Transaction('buy', '2023-08-09', 5, 49.7, -248.48),
    Transaction('sell', '2023-08-16', -8, 49.7, 397.63),
    Transaction('sell', '2023-08-17', -8, 49.51, 396.11),
    Transaction('sell', '2023-08-18', -6, 49.46, 296.78),
    Transaction('buy', '2023-08-28', 5, 49.16, -245.79),
    Transaction('buy', '2023-09-01', 5, 49.27, -246.35),
    Transaction('buy', '2023-09-05', 2, 49.2, -98.40),
    Transaction('buy', '2023-09-06', 3, 49.16, -147.48),
    Transaction('buy', '2023-09-08', 1, 49.2, -49.20),
    Transaction('buy', '2023-09-15', 2, 49.08, -98.16),
    Transaction('sell', '2023-10-03', -5, 47.61, 238.06),
    Transaction('sell', '2023-10-16', -1358, 47.96, 65129.82),
    Transaction('sell', '2023-10-20', -6, 47.4, 284.40),
    Transaction('sell', '2023-10-23', -3, 47.44, 142.32),
    Transaction('sell', '2023-10-25', -2, 47.38, 94.76),
    Transaction('sell', '2023-10-26', -18, 47.45, 854.15),
    Transaction('sell', '2023-10-27', -1, 47.47, 47.47),
    Transaction('sell', '2023-11-08', -191, 48.56, 9275.69),
    Transaction('sell', '2023-11-13', -3, 48.52, 145.56),
    Transaction('buy', '2023-12-22', 3, 50.92, -152.78),
    Transaction('buy', '2023-12-22', 1, 50.93, -50.93),
    Transaction('buy', '2023-12-27', 2, 51.06, -102.11),
    Transaction('buy', '2023-12-27', 1, 51.06, -51.06),
    Transaction('buy', '2023-12-28', 23, 51.04, -1173.81),
    Transaction('buy', '2024-01-05', 4, 50.9, -203.59),
    Transaction('buy', '2024-01-05', 244, 50.9, -12418.38),
    Transaction('buy', '2024-01-26', 3, 50.5, -151.50),
    Transaction('buy', '2024-01-26', 1, 50.5, -50.50),
    Transaction('buy', '2024-02-07', 6, 50.63, -303.78),
    Transaction('buy', '2024-02-07', 1, 50.63, -50.63),
    Transaction('buy', '2024-02-08', 1, 50.56, -50.57),
    Transaction('buy', '2024-02-12', 1, 50.71, -50.71),
    Transaction('buy', '2024-02-12', 1, 50.71, -50.71),
    Transaction('buy', '2024-02-21', 5, 50.74, -253.68),
    Transaction('buy', '2024-02-21', 1, 50.74, -50.74),
    Transaction('buy', '2024-03-07', 1, 50.78, -50.78),
    Transaction('buy', '2024-03-08', 3, 50.8, -152.40),
    Transaction('buy', '2024-03-08', 1, 50.8, -50.80),
    Transaction('buy', '2024-03-11', 2, 50.8, -101.60),
    Transaction('buy', '2024-03-11', 1, 50.8, -50.80),
    Transaction('buy', '2024-03-12', 1, 50.74, -50.74),
    Transaction('buy', '2024-03-21', 1, 50.62, -50.62),
    Transaction('buy', '2024-03-21', 1, 50.62, -50.62),
    Transaction('buy', '2024-03-22', 1, 50.75, -50.75),
    Transaction('buy', '2024-03-22', 3, 50.75, -152.25),
    Transaction('buy', '2024-03-26', 1, 50.51, -50.51),
    Transaction('buy', '2024-03-28', 6, 50.54, -303.27),
    Transaction('buy', '2024-03-28', 1, 50.54, -50.55),
    Transaction('buy', '2024-04-01', 12, 50.26, -603.18),
    Transaction('buy', '2024-04-01', 1, 50.27, -50.27),
    Transaction('buy', '2024-04-02', 6, 50.2, -301.17),
    Transaction('buy', '2024-04-02', 1, 50.2, -50.20),
    Transaction('buy', '2024-04-05', 1, 50.02, -50.02),
    Transaction('buy', '2024-04-08', 1, 50.03, -50.03),
    Transaction('buy', '2024-04-08', 3, 50.02, -150.08),
    Transaction('buy', '2024-04-10', 1, 49.89, -49.89),
    Transaction('buy', '2024-04-12', 3, 50.06, -150.20),
    Transaction('buy', '2024-04-12', 1, 50.06, -50.07),
    Transaction('buy', '2024-04-25', 8, 49.75, -397.97),
    Transaction('buy', '2024-04-25', 1, 49.74, -49.75),
    Transaction('buy', '2024-05-08', 2, 50.28, -100.55),
    Transaction('buy', '2024-05-08', 1, 50.28, -50.28),
    Transaction('buy', '2024-05-13', 1, 50.22, -50.22),
    Transaction('buy', '2024-05-21', 1, 50.08, -50.08),
    Transaction('buy', '2024-05-21', 1, 50.08, -50.08),
    Transaction('buy', '2024-05-22', 5, 49.95, -249.74),
    Transaction('buy', '2024-05-22', 1, 49.95, -49.95),
    Transaction('buy', '2024-05-29', 2, 49.58, -99.15),
    Transaction('buy', '2024-06-10', 3, 49.9, -149.72),
    Transaction('buy', '2024-06-10', 1, 49.91, -49.91),
    Transaction('buy', '2024-06-14', 1, 50.36, -50.36),
    Transaction('buy', '2024-06-18', 11, 50.33, -553.63),
    Transaction('buy', '2024-06-18', 1, 50.33, -50.33),
    Transaction('buy', '2024-06-20', 1, 50.32, -50.33),
    Transaction('buy', '2024-06-26', 3, 50.2, -150.62),
    Transaction('buy', '2024-06-26', 1, 50.21, -50.21),
    Transaction('buy', '2024-06-27', 1, 50.24, -50.24),
    Transaction('buy', '2024-06-27', 2, 50.24, -100.47),
    Transaction('buy', '2024-07-02', 10, 50.08, -500.85),
    Transaction('buy', '2024-07-02', 1, 50.08, -50.09),
    Transaction('buy', '2024-07-05', 7, 50.29, -352.01),
    Transaction('buy', '2024-07-05', 1, 50.29, -50.29),
    Transaction('sell', '2024-08-07', -278, 50.69, 14091.85),
    Transaction('buy', '2024-09-09', 3, 50.89, -152.67),
    Transaction('buy', '2024-09-09', 1, 50.88, -50.89),
    Transaction('buy', '2024-09-17', 1, 51.07, -51.07),
    Transaction('buy', '2024-09-18', 1, 51.05, -51.05))


# This is the actual Wealthfront cost basis. We're trying
# to pick a strategy that chooses lots to sell that ends
# up with these lots and holdings at the end.
actual_cost_basis = frozenset((
    ('2023-05-19', 1, '50.30'),
    ('2023-05-19', 4, '199.10'),
    ('2023-06-02', 5, '249.30'),
    ('2023-06-05', 4, '199.60'),
    ('2023-06-07', 1, '50.00'),
    ('2023-06-12', 1, '49.93'),
    ('2023-06-14', 10, '500.70'),
    ('2023-06-16', 12, '601.92'),
    ('2023-07-10', 1, '49.93'),
    ('2023-07-11', 2, '99.94'),
    ('2023-08-01', 3, '149.85'),
    ('2023-08-02', 5, '248.88'),
    ('2023-08-03', 2, '100.45'),
    ('2023-08-03', 2, '99.95'),
    ('2023-08-07', 3, '150.76'),
    ('2023-08-28', 5, '245.79'),
    ('2023-09-01', 5, '246.35'),
    ('2023-09-05', 2, '98.40'),
    ('2023-09-06', 3, '147.48'),
    ('2023-09-08', 1, '49.20'),
    ('2023-09-15', 2, '98.16'),
    ('2024-01-26', 1, '50.50'),
    ('2024-01-26', 3, '151.50'),
    ('2024-02-07', 1, '50.63'),
    ('2024-02-07', 6, '303.78'),
    ('2024-02-08', 1, '50.57'),
    ('2024-02-12', 1, '50.71'),
    ('2024-02-21', 1, '50.74'),
    ('2024-02-21', 5, '253.68'),
    ('2024-03-07', 1, '50.78'),
    ('2024-03-08', 1, '50.80'),
    ('2024-03-08', 3, '152.40'),
    ('2024-03-11', 1, '50.80'),
    ('2024-03-11', 2, '101.60'),
    ('2024-03-12', 1, '50.74'),
    ('2024-03-21', 1, '50.62'),
    ('2024-03-22', 1, '50.75'),
    ('2024-03-22', 3, '152.25'),
    ('2024-03-26', 1, '50.51'),
    ('2024-03-28', 1, '50.55'),
    ('2024-03-28', 6, '303.27'),
    ('2024-04-01', 1, '50.27'),
    ('2024-04-01', 12, '603.18'),
    ('2024-04-02', 1, '50.20'),
    ('2024-04-02', 6, '301.17'),
    ('2024-04-05', 1, '50.02'),
    ('2024-04-08', 1, '50.03'),
    ('2024-04-08', 3, '150.08'),
    ('2024-04-10', 1, '49.89'),
    ('2024-04-12', 1, '50.07'),
    ('2024-04-12', 3, '150.20'),
    ('2024-04-25', 1, '49.75'),
    ('2024-04-25', 8, '397.97'),
    ('2024-05-08', 1, '50.28'),
    ('2024-05-08', 2, '100.55'),
    ('2024-05-13', 1, '50.22'),
    ('2024-05-21', 1, '50.08'),
    ('2024-05-22', 1, '49.95'),
    ('2024-05-22', 5, '249.74'),
    ('2024-05-29', 2, '99.15'),
    ('2024-06-10', 1, '49.91'),
    ('2024-06-10', 3, '149.72'),
    ('2024-06-14', 1, '50.36'),
    ('2024-06-18', 1, '50.33'),
    ('2024-06-18', 11, '553.63'),
    ('2024-06-20', 1, '50.33'),
    ('2024-06-26', 1, '50.21'),
    ('2024-06-26', 3, '150.62'),
    ('2024-06-27', 1, '50.24'),
    ('2024-06-27', 2, '100.47'),
    ('2024-07-02', 1, '50.09'),
    ('2024-07-02', 10, '500.85'),
    ('2024-07-05', 1, '50.29'),
    ('2024-07-05', 7, '352.01'),
    ('2024-09-09', 1, '50.89'),
    ('2024-09-09', 3, '152.67'),
    ('2024-09-17', 1, '51.07'),
    ('2024-09-18', 1, '51.05')))

# Let's make it easier for now, don't worry about the price,
# If we can match a set of lots based on date and shares,
# Eg., ('2024-09-18', 1)
# Then that's a "win" condition.
acb_date_and_shares = set([i[:2] for i in actual_cost_basis])


def lifo_picker(lots: Sequence[Lot], date: str) -> int:
    """LIFO: Last In First Out. Easy, last one is index -1. """
    return -1


def fifo_picker(lots: Sequence[Lot], date: str) -> int:
    """FIFO: First In First Out. Easy, first one is index 0."""
    return 0


def hifo_picker(lots: Sequence[Lot], date: str) -> int:
    """HIFO: Highest (unit price) In First Out"""
    max_index, max_value = max(enumerate(lots), key=lambda x: x[1].unit_price)
    return max_index


def best_matching_picker(lots: Sequence[Lot], date: str) -> int:
    """Try to match actual Cost Basis"""
    # Best factors are 1.0156 and 1.0192
    long_term_benefit_factor = 1.0156

    max_t = -1
    year_before = f'{int(date[:4]) - 1}{date[4:]}'
    iso_date = datetime.date.fromisoformat(date[:10])
    for i, t in enumerate(lots):
        adj_price = t.unit_price
        if t.date <= year_before:
            adj_price *= long_term_benefit_factor
        if adj_price > max_t:
            # 61 days to avoid a wash sale
            if (iso_date - datetime.date.fromisoformat(t.date[:10])).days >= 61:
                max_i = i
                max_t = adj_price
    return max_i


def compare_lots_to_actual(lots: Sequence[Lot]) -> None:
    """We'll make our comparisons by doing a set intersection of our
    final lots vs. the actual cost basis lots.

    For example:
    /‾‾‾‾‾‾‾‾‾‾‾/‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾\‾‾‾‾‾‾‾‾‾‾‾\
    ⎸  ours  3  ⎸ intersection 60 ⎹ actual 17 ⎹
    \___________\________________/___________/

    Means 60 of the lots are in both sets. Our attempt has 3 lots not
    in the actual cost basis. And the actual cost basis has 17 lots that
    are not in ours.

    The winning condition where everything matches is:
    /‾‾‾‾‾‾‾‾‾‾‾/‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾\‾‾‾‾‾‾‾‾‾‾‾\
    ⎸  ours  0  ⎸ intersection 78 ⎹ actual  0 ⎹
    \___________\________________/___________/
    """
    ours = set()
    for t in lots:
        ours.add((t.date, t.units))
    only_ours = len(ours - acb_date_and_shares)
    common = ours & acb_date_and_shares
    only_actual = len(acb_date_and_shares - ours)
    print('/‾‾‾‾‾‾‾‾‾‾‾/‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾\\‾‾‾‾‾‾‾‾‾‾‾\\')
    print(f'⎸  ours {only_ours:>2}  ⎸ intersection {len(common):>2} ⎹ actual {only_actual:>2} ⎹')
    print('\\___________\\_________________/___________/')


def try_strategy(select_lot: Callable) -> None:
    lots = []
    for t in transactions:
        if t.type == 'buy':
            lots.append(Lot(t.trade_date, t.units, t.unit_price))
        else:  # sell
            units = -t.units
            while units:
                lot_id = select_lot(lots, t.trade_date)
                lot = lots[lot_id]
                if units >= lot.units:
                    units -= lot.units
                    lots.pop(lot_id)
                else:
                    lot.units -= units
                    units = 0
    units_sum = sum([t.units for t in lots])
    price_sum = sum([t.unit_price * t.units for t in lots])
    print(f"{select_lot.__name__.upper()}:")
    compare_lots_to_actual(lots)
    print()


if __name__ == '__main__':
    try_strategy(lifo_picker)
    try_strategy(fifo_picker)
    try_strategy(hifo_picker)
    try_strategy(best_matching_picker)