@startuml

set namespaceseparator none
skinparam dpi 96
skinparam ArrowColor Black
skinparam shadowing true
skinparam class {
  backgroundColor Business
  borderColor Red
}

package pytorch #DDDDDD {}
package tket #DDDDDD {}
package pennylane #DDDDDD {}

abstract class Model {
    symbols
    weights
    {abstract} initialise_weights()
    {abstract}{static} from_checkpoint()
    {abstract} get_diagram_output()
    {abstract} forward()
    {static} from diagrams()
}

abstract class Trainer {
    backend
    model
    load_training_checkpoint()
    save_checkpoint()
    {abstract} training_step()
    {abstract} validation_step()
    fit()
}

abstract class Optimizer {
    model
    loss
    hyperparams: dict[str, float]
    bounds
    {abstract} backward()
    {abstract} step()
    {abstract} state_dict()
    {abstract} load_state_dict()
    zero_grad()
}

class Dataset {
    data
    targets
    batch_size
    shuffle
    {static} shuffle_data()
}

class CheckPoint {
    entries
    add_many()
    {static} from_file()
    to_file()
}

class PytorchModel {}
class PennyLaneModel {}
class NumpyModel {
    use_jit
    lambdas
}
class TketModel {}
class QuantumModel {
    {static} SMOOTHING
}

class QuantumTrainer {
    optimizer
}

class PytorchTrainer {
    optimizer
}

class SPSAOptimizer {}
class NelderMeadOptimizer {}
class RotosolveOptimizer {}

class BinaryCrossEntropyLoss {}

class CrossEntropyLoss {
  calculate_loss()
}

class LossFunction {
  backend : module
  {abstract}calculate_loss()
}

class MSELoss {
  calculate_loss()
}

Model <|-- PytorchModel
Model <|-- PennyLaneModel
Model <|-- QuantumModel
QuantumModel <|-- TketModel
QuantumModel <|-- NumpyModel
Trainer <|-- PytorchTrainer
Trainer <|-- QuantumTrainer

Optimizer <|-- SPSAOptimizer
Optimizer <|-- NelderMeadOptimizer
Optimizer <|-- RotosolveOptimizer

PennyLaneModel -- pennylane
PennyLaneModel -- pytorch
PytorchModel -- pytorch
PytorchTrainer -- pytorch
TketModel -- tket
QuantumTrainer -- tket

Trainer *-- CheckPoint
CheckPoint --* Model
Trainer::model *- Model
Trainer *-u- Dataset
QuantumModel -* Optimizer::model
Optimizer -* QuantumTrainer::optimizer
LossFunction --* Optimizer::loss

PennyLaneModel -- PytorchTrainer: usedWith
PytorchModel -- PytorchTrainer: usedWith
NumpyModel -- QuantumTrainer: usedWith
TketModel -- QuantumTrainer: usedWith
QuantumModel -- LossFunction: uses

CrossEntropyLoss <|-- BinaryCrossEntropyLoss
LossFunction <|-- CrossEntropyLoss
LossFunction <|-- MSELoss

@enduml
