KinderElektro/src/main/java/de/ph87/electro/circuit/calculation/Calculation.java

package de.ph87.electro.circuit.calculation;

import de.ph87.electro.CONFIG;
import de.ph87.electro.circuit.Circuit;
import de.ph87.electro.circuit.part.InnerConnection;
import de.ph87.electro.circuit.part.Part;
import de.ph87.electro.circuit.part.node.Node;
import de.ph87.electro.circuit.part.parts.Battery;
import de.ph87.electro.circuit.wire.Wire;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.math3.linear.*;

import java.util.*;

import static java.lang.Math.max;

@Slf4j
@Getter
public class Calculation {

  private final List<Node> nodes = new ArrayList<>();

  private final Set<Part> parts = new HashSet<>();

  private final Set<Wire> wires = new HashSet<>();

  private final RealMatrix matrix;

  private final RealVector currents;

  private RealVector potentials = null;

  public Calculation(final int numNodes) {
    matrix = new Array2DRowRealMatrix(numNodes, numNodes);
    currents = new ArrayRealVector(numNodes);
  }

  private Calculation(final Set<Node> connectedNodes, final Battery pivot) {
    for (final Node node : connectedNodes) {
      if (node instanceof final Node partNode) {
        parts.add(partNode.getPart());
      }
      if (!nodes.contains(node) && node != pivot.getMinus()) {
        // pivot.minus is GND and cannot be part of the matrix (linear dependency)
        this.nodes.add(node);
      }
      wires.addAll(node.getWires());
    }

    matrix = new Array2DRowRealMatrix(this.nodes.size(), this.nodes.size());
    currents = new ArrayRealVector(this.nodes.size());
    fromSchematic();
    solve();
    toSchematic();
  }

  public static List<Calculation> calculate(final Circuit circuit) {
    final List<Calculation> calculations = new ArrayList<>();
    final List<Battery> batteries = new ArrayList<>(circuit.streamParts().flatMap(Battery::filterCast).toList());
    while (!batteries.isEmpty()) {
      final Battery pivot = batteries.removeFirst();
      final Set<Node> connectedNodes = new HashSet<>();
      pivot.getPlus().collectConnectedNodes(connectedNodes);
      pivot.getMinus().collectConnectedNodes(connectedNodes);
      connectedNodes.stream().map(Node::getPart).flatMap(Battery::filterCast).forEach(batteries::remove);
      calculations.add(new Calculation(connectedNodes, pivot));
    }
    return calculations;
  }

  private void fromSchematic() {
    wires.forEach(
      wire -> addResistor(wire.getA(), wire.getB(), CONFIG.NO_RESISTANCE)
    );
    parts.forEach(part -> {
      for (final InnerConnection innerConnection : part.getInnerConnections()) {
        addResistor(innerConnection.a, innerConnection.b, innerConnection.resistance);
      }
      if (part instanceof final Battery battery) {
        final double current = battery.getVoltage() / battery.getResistance();
        final int indexMinus = getNodeIndex(battery.getMinus());
        final int indexPlus = getNodeIndex(battery.getPlus());
        addCurrentSource(indexMinus, indexPlus, current);
      }
    });
  }

  private void toSchematic() {
    wires.forEach(wire -> wire.setCurrent(getCurrent(wire)));
    parts.forEach(part -> {
      part.getNodes().forEach(node -> node.setVoltage(getPotential(node)));
      part.postCalculate();
    });
  }

  private double getPotential(final @NonNull Node node) {
    final int index = getNodeIndex(node);
    if (index < 0) {
      return 0; // per definition
    }
    if (potentials == null) {
      return Double.NaN;
    }
    return potentials.getEntry(index);
  }

  private double getCurrent(final Wire wire) {
    final int indexA = getNodeIndex(wire.getA());
    final int indexB = getNodeIndex(wire.getB());
    final double conductance = indexA >= 0 && indexB >= 0 ? matrix.getEntry(indexA, indexB) : 1 / CONFIG.NO_RESISTANCE;
    final double potentialDifference = getPotential(wire.getB()) - getPotential(wire.getA());
    return conductance * potentialDifference;
  }

  private int getNodeIndex(@NonNull final Node node) {
    for (int i = 0; i < nodes.size(); i++) {
      if (nodes.get(i) == node) {
        return i;
      }
    }
    return -1;
  }

  public void addResistor(final Node a, final Node b, final double resistance) {
    final int indexA = getNodeIndex(a);
    final int indexB = getNodeIndex(b);
    addResistor(indexA, indexB, resistance);
  }

  public void addResistor(final int index0, final int index1, double resistance) {
    double conductance = 1.0 / resistance;
    if (index0 >= 0) {
      matrix.setEntry(index0, index0, matrix.getEntry(index0, index0) + conductance);
    }
    if (index1 >= 0) {
      matrix.setEntry(index1, index1, matrix.getEntry(index1, index1) + conductance);
    }
    if (index0 >= 0 && index1 >= 0) {
      matrix.setEntry(index0, index1, matrix.getEntry(index0, index1) - conductance);
      matrix.setEntry(index1, index0, matrix.getEntry(index1, index0) - conductance);
    }
    potentials = null;
  }

  public void addCurrentSource(final int index0, final int index1, final double current) {
    if (index0 >= 0) {
      currents.setEntry(index0, currents.getEntry(index0) - current);
    }
    if (index1 >= 0) {
      currents.setEntry(index1, currents.getEntry(index1) + current);
    }
    potentials = null;
  }

  public void solve() {
    final DecompositionSolver solver = new LUDecomposition(matrix).getSolver();
    try {
      potentials = solver.solve(currents);
    } catch (SingularMatrixException e) {
      potentials = null;
      log.error("Schaltung fehlerhaft: {}", e.getMessage());
    }
  }

  @Override
  public String toString() {
    final Format matrixFormat = new Format(matrix);
    final Format currentsFormat = new Format(currents);
    final Format potentialsFormat = new Format(potentials);

    final StringBuilder builder = new StringBuilder();
    for (int row = 0; row < matrix.getRowDimension(); row++) {
      builder.append("[");
      String prefix = "";
      for (int c = 0; c < matrix.getColumnDimension(); c++) {
        builder.append(prefix);
        builder.append(matrixFormat.format(matrix.getEntry(row, c)));
        prefix = ", ";
      }

      final String potentialString = potentialsFormat.format(potentials.getEntry(row));
      final String currentString = currentsFormat.format(currents.getEntry(row));
      builder.append("] x [%s V] = [%s A]\n".formatted(potentialString, currentString));
    }
    return builder.toString();
  }

  @Getter
  private static final class Format {

    private final String format;

    private int integer;

    private int decimal;

    public Format(final RealMatrix matrix) {
      for (int r = 0; r < matrix.getRowDimension(); r++) {
        for (int c = 0; c < matrix.getColumnDimension(); c++) {
          append(matrix.getEntry(r, c));
        }
      }
      format = calculateFormat();
    }

    public Format(final RealVector vector) {
      for (int r = 0; r < vector.getDimension(); r++) {
        append(vector.getEntry(r));
      }
      format = calculateFormat();
    }

    private void append(final double value) {
      final String string = "%+f".formatted(value);
      final String[] parts = string.split("[.,]");
      integer = max(integer, parts[0].length());
      if (parts.length > 1) {
        decimal = max(decimal, parts[1].length());
      }
    }

    private String calculateFormat() {
      return "%%+%d.%df".formatted(integer + decimal + (decimal > 0 ? 1 : 0), decimal);
    }

    public String format(final double value) {
      return format.formatted(value);
    }

  }

}