fisk/core/models.py

from datetime import timedelta

import stripe
from django.conf import settings
from django.contrib.auth.models import AbstractUser
from django.db import models

from core.exchanges.alpaca import AlpacaExchange
from core.exchanges.oanda import OANDAExchange
from core.lib.customers import get_or_create, update_customer_fields
from core.util import logs

log = logs.get_logger(__name__)
EXCHANGE_MAP = {"alpaca": AlpacaExchange, "oanda": OANDAExchange}
TYPE_CHOICES = (
    ("market", "Market"),
    ("limit", "Limit"),
)
DIRECTION_CHOICES = (
    ("buy", "Buy"),
    ("sell", "Sell"),
)
TIF_CHOICES = (
    ("gtc", "GTC (Good Til Cancelled)"),
    ("gfd", "GFD (Good For Day)"),
    ("fok", "FOK (Fill Or Kill)"),
    ("ioc", "IOC (Immediate Or Cancel)"),
)
DAY_CHOICES = (
    (1, "Monday"),
    (2, "Tuesday"),
    (3, "Wednesday"),
    (4, "Thursday"),
    (5, "Friday"),
    (6, "Saturday"),
    (7, "Sunday"),
)
SIGNAL_TYPE_CHOICES = (
    ("entry", "Entry"),
    ("exit", "Exit"),
    ("trend", "Trend"),
)


class Plan(models.Model):
    name = models.CharField(max_length=255, unique=True)
    description = models.CharField(max_length=1024, null=True, blank=True)
    cost = models.IntegerField()
    product_id = models.CharField(max_length=255, unique=True, null=True, blank=True)
    image = models.CharField(max_length=1024, null=True, blank=True)

    def __str__(self):
        return f"{self.name} (£{self.cost})"


class User(AbstractUser):
    # Stripe customer ID
    stripe_id = models.CharField(max_length=255, null=True, blank=True)
    last_payment = models.DateTimeField(null=True, blank=True)
    plans = models.ManyToManyField(Plan, blank=True)
    email = models.EmailField(unique=True)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._original = self

    def save(self, *args, **kwargs):
        """
        Override the save function to create a Stripe customer.
        """
        if settings.STRIPE_ENABLED:
            if not self.stripe_id:  # stripe ID not stored
                self.stripe_id = get_or_create(
                    self.email, self.first_name, self.last_name
                )

            to_update = {}
            if self.email != self._original.email:
                to_update["email"] = self.email
            if self.first_name != self._original.first_name:
                to_update["first_name"] = self.first_name
            if self.last_name != self._original.last_name:
                to_update["last_name"] = self.last_name

            update_customer_fields(self.stripe_id, **to_update)

        super().save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        if settings.STRIPE_ENABLED:
            if self.stripe_id:
                stripe.Customer.delete(self.stripe_id)
                log.info(f"Deleted Stripe customer {self.stripe_id}")
        super().delete(*args, **kwargs)

    def has_plan(self, plan):
        plan_list = [plan.name for plan in self.plans.all()]
        return plan in plan_list

    def get_notification_settings(self):
        return NotificationSettings.objects.get_or_create(user=self)[0]


class Account(models.Model):
    EXCHANGE_CHOICES = (("alpaca", "Alpaca"), ("oanda", "OANDA"))
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    exchange = models.CharField(choices=EXCHANGE_CHOICES, max_length=255)
    api_key = models.CharField(max_length=255)
    api_secret = models.CharField(max_length=255)
    sandbox = models.BooleanField(default=False)
    enabled = models.BooleanField(default=True)
    supported_symbols = models.JSONField(default=list)
    instruments = models.JSONField(default=list)
    currency = models.CharField(max_length=255, null=True, blank=True)
    risk_model = models.ForeignKey(
        "core.RiskModel", on_delete=models.SET_NULL, null=True, blank=True
    )
    initial_balance = models.FloatField(default=0)

    def __str__(self):
        name = f"{self.name} ({self.exchange})"
        if self.sandbox:
            name += " (sandbox)"
        return name

    def update_info(self, save=True):
        client = self.get_client()
        if client:
            response = client.get_instruments()
            supported_symbols = client.get_supported_assets(response)
            currency = client.get_account()["currency"]
            log.debug(f"Supported symbols for {self.name}: {supported_symbols}")
            self.supported_symbols = supported_symbols
            self.instruments = response
            self.currency = currency
        if save:
            self.save()

    def save(self, *args, **kwargs):
        """
        Override the save function to update supported symbols.
        """
        self.update_info(save=False)
        super().save(*args, **kwargs)

    def get_client(self):
        if self.exchange in EXCHANGE_MAP:
            return EXCHANGE_MAP[self.exchange](self)
        else:
            raise Exception(f"Exchange not supported : {self.exchange}")

    @property
    def client(self):
        """
        Convenience property for one-off API calls.
        """
        return self.get_client()

    @property
    def rawclient(self):
        """
        Convenience property for one-off API calls.
        """
        return self.get_client().client

    @classmethod
    def get_by_id(cls, account_id, user):
        return cls.objects.get(id=account_id, user=user)

    @classmethod
    def get_by_id_no_user_check(cls, account_id):
        return cls.objects.get(id=account_id)


class Session(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    request = models.CharField(max_length=255, null=True, blank=True)
    session = models.CharField(max_length=255, null=True, blank=True)
    subscription_id = models.CharField(max_length=255, null=True, blank=True)
    plan = models.ForeignKey(Plan, null=True, blank=True, on_delete=models.CASCADE)


class Hook(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=1024)
    hook = models.CharField(max_length=255, unique=True)  # hook URL
    received = models.IntegerField(default=0)

    def __str__(self):
        return f"{self.name} ({self.hook})"


class Signal(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=1024)
    signal = models.CharField(max_length=256)  # signal name
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE)
    direction = models.CharField(choices=DIRECTION_CHOICES, max_length=255)
    received = models.IntegerField(default=0)
    type = models.CharField(choices=SIGNAL_TYPE_CHOICES, max_length=255)

    def __str__(self):
        return f"{self.name} ({self.hook.name}) - {self.direction}"


class Trade(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE, null=True, blank=True)
    signal = models.ForeignKey(Signal, on_delete=models.CASCADE, null=True, blank=True)
    symbol = models.CharField(max_length=255)
    time_in_force = models.CharField(choices=TIF_CHOICES, max_length=255, default="gtc")
    type = models.CharField(choices=TYPE_CHOICES, max_length=255)
    amount = models.FloatField(null=True, blank=True)
    amount_usd = models.FloatField(null=True, blank=True)
    price = models.FloatField(null=True, blank=True)
    stop_loss = models.FloatField(null=True, blank=True)
    trailing_stop_loss = models.FloatField(null=True, blank=True)
    take_profit = models.FloatField(null=True, blank=True)
    status = models.CharField(max_length=255, null=True, blank=True)
    information = models.JSONField(null=True, blank=True)
    direction = models.CharField(choices=DIRECTION_CHOICES, max_length=255)

    # To populate from the trade
    order_id = models.CharField(max_length=255, null=True, blank=True)
    client_order_id = models.CharField(max_length=255, null=True, blank=True)
    response = models.JSONField(null=True, blank=True)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._original = self

    def post(self):
        if self.status in ["rejected", "close"]:
            log.debug(f"Trade {self.id} rejected. Not posting.")
            log.debug(f"Trade {self.id} information: {self.information}")
        else:
            return self.account.client.post_trade(self)

    def delete(self, *args, **kwargs):
        # close the trade
        super().delete(*args, **kwargs)

    @classmethod
    def get_by_id(cls, trade_id, user):
        return cls.objects.get(id=trade_id, user=user)

    @classmethod
    def get_by_id_or_order(cls, trade_id, account_id, user):
        try:
            account = Account.objects.get(id=account_id, user=user)
        except Account.DoesNotExist:
            return None
        try:
            return cls.objects.get(id=trade_id, account=account, user=user)
        except cls.DoesNotExist:
            try:
                return cls.objects.get(order_id=trade_id, account=account, user=user)
            except cls.DoesNotExist:
                return None


class Callback(models.Model):
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE)
    signal = models.ForeignKey(Signal, on_delete=models.CASCADE)
    title = models.CharField(max_length=1024, null=True, blank=True)
    message = models.CharField(max_length=1024, null=True, blank=True)
    period = models.CharField(max_length=255, null=True, blank=True)
    sent = models.BigIntegerField(null=True, blank=True)
    trade = models.BigIntegerField(null=True, blank=True)
    exchange = models.CharField(max_length=255, null=True, blank=True)
    base = models.CharField(max_length=255, null=True, blank=True)
    quote = models.CharField(max_length=255, null=True, blank=True)
    contract = models.CharField(max_length=255, null=True, blank=True)
    price = models.FloatField(null=True, blank=True)
    symbol = models.CharField(max_length=255)


class TradingTime(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    description = models.TextField(null=True, blank=True)
    start_day = models.IntegerField(choices=DAY_CHOICES)
    end_day = models.IntegerField(choices=DAY_CHOICES)
    start_time = models.TimeField()
    end_time = models.TimeField()

    def within_range(self, ts):
        """
        Check if the specified time is within the configured trading times.
        :param ts: Timestamp
        :type ts: datetime
        :return: whether or not the time is within the trading range
        :rtype: bool
        """
        start_day = self.start_day
        end_day = self.end_day
        # Check the day is between the start and end day
        if not start_day <= ts.weekday() + 1 <= end_day:
            return False

        start_time = self.start_time
        end_time = self.end_time

        # Get what the start time would be this week
        ts_monday = ts - timedelta(days=ts.weekday())

        # Now we need to add our day of week to monday
        # Let's set the offset now since it's off by one
        offset_start = start_day - 1
        # Datetime: monday=0, tuesday=1, us: monday=1, tuesday=2, so we need to subtract
        # one from ours to not be off by one
        offset_end = end_day - 1

        # Now we can add the offset to the monday
        start = ts_monday + timedelta(days=offset_start)
        start = start.replace(
            hour=start_time.hour,
            minute=start_time.minute,
            second=start_time.second,
            microsecond=start_time.microsecond,
        )
        end = ts_monday + timedelta(days=offset_end)
        end = end.replace(
            hour=end_time.hour,
            minute=end_time.minute,
            second=end_time.second,
            microsecond=end_time.microsecond,
        )
        # Check if the ts is between the start and end times
        # ts must be more than start and less than end
        return ts >= start and ts <= end
        return True

    def __str__(self):
        return (
            f"{self.name} ({self.get_start_day_display()} at {self.start_time} - "
            f"{self.get_end_day_display()} at {self.end_time})"
        )


class Strategy(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    description = models.TextField(null=True, blank=True)
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    trading_times = models.ManyToManyField(TradingTime)
    order_type = models.CharField(
        choices=TYPE_CHOICES, max_length=255, default="market"
    )
    time_in_force = models.CharField(choices=TIF_CHOICES, max_length=255, default="gtc")
    entry_signals = models.ManyToManyField(
        Signal, related_name="entry_strategies", blank=True
    )
    exit_signals = models.ManyToManyField(
        Signal, related_name="exit_strategies", blank=True
    )
    trend_signals = models.ManyToManyField(
        Signal, related_name="trend_strategies", blank=True
    )
    enabled = models.BooleanField(default=False)
    take_profit_percent = models.FloatField(default=1.5)
    stop_loss_percent = models.FloatField(default=1.0)
    trailing_stop_loss_percent = models.FloatField(default=1.0, null=True, blank=True)
    price_slippage_percent = models.FloatField(default=2.5)
    callback_price_deviation_percent = models.FloatField(default=0.5)
    trade_size_percent = models.FloatField(default=0.5)
    trends = models.JSONField(null=True, blank=True)

    class Meta:
        verbose_name_plural = "strategies"

    def __str__(self):
        return self.name


class NotificationSettings(models.Model):
    user = models.OneToOneField(User, on_delete=models.CASCADE)
    ntfy_topic = models.CharField(max_length=255, null=True, blank=True)
    ntfy_url = models.CharField(max_length=255, null=True, blank=True)

    def __str__(self):
        return f"Notification settings for {self.user}"


class RiskModel(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    description = models.TextField(null=True, blank=True)
    # Maximum amount of money to have lost from the initial balance to stop trading
    max_loss_percent = models.FloatField(default=0.05)
    # Maximum amount of money to risk on all open trades
    max_risk_percent = models.FloatField(default=0.05)
    # Maximum number of trades
    max_open_trades = models.IntegerField(default=10)
    # Maximum number of trades per symbol
    max_open_trades_per_symbol = models.IntegerField(default=2)

    def __str__(self):
        return self.name


class AssetGroup(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    description = models.TextField(null=True, blank=True)

    # Account for checking pairs on children if specified
    account = models.ForeignKey(Account, on_delete=models.CASCADE)

    # Dict like {"RUB": True, "USD": False}
    allowed = models.JSONField(null=True, blank=True, default=dict)

    def __str__(self):
        return self.name

    @property
    def matches(self):
        """
        Get the total number of matches for this group.
        """
        if isinstance(self.allowed, dict):
            truthy_values = [x for x in self.allowed.values() if x is True]
            return f"{len(truthy_values)}/{len(self.allowed)}"

    @property
    def restrictions(self):
        """
        Get the total number of restrictions for this group.
        """
        return self.assetrestriction_set.count()


class AssetRestriction(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    description = models.TextField(null=True, blank=True)
    pairs = models.CharField(max_length=4096, null=True, blank=True)
    pairs_parsed = models.JSONField(null=True, blank=True)

    group = models.ForeignKey(
        AssetGroup, on_delete=models.CASCADE, null=True, blank=True
    )