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

package de.ph87.electro.circuit.calculation;

import de.ph87.electro.circuit.Circuit;
import de.ph87.electro.circuit.wire.Wire;
import de.ph87.electro.circuit.part.InnerConnection;
import de.ph87.electro.circuit.part.junction.Junction;
import de.ph87.electro.circuit.part.Part;
import de.ph87.electro.circuit.part.parts.PartBattery;
import de.ph87.electro.circuit.part.parts.PartLight;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.math3.linear.*;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import static java.lang.Math.max;

@Slf4j
@Getter
public class Calculation {

  public static final double NO_RESISTANCE = 1e-12;

  private static final double FULL_ADMITTANCE = 1 / NO_RESISTANCE;

  private final List<Junction> junctions = 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 static List<Calculation> calculate(final Circuit circuit) {
    final List<Calculation> calculations = new ArrayList<>();
    final List<PartBattery> batteries = new ArrayList<>(circuit.streamParts().flatMap(PartBattery::filterCast).toList());
    while (!batteries.isEmpty()) {
      final PartBattery pivot = batteries.removeFirst();
      final Set<Junction> connectedJunctions = new HashSet<>();
      pivot.getPlus().collectConnectedJunctions(connectedJunctions);
      pivot.getMinus().collectConnectedJunctions(connectedJunctions);
      connectedJunctions.stream().map(Junction::getOwner).flatMap(PartBattery::filterCast).forEach(batteries::remove);
      calculations.add(new Calculation(connectedJunctions, pivot));
    }
    return calculations;
  }

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

  private Calculation(final Set<Junction> connectedJunctions, final PartBattery pivot) {
    for (final Junction junction : connectedJunctions) {
      parts.add(junction.getOwner());
      if (!junctions.contains(junction) && junction != pivot.getMinus()) {
        // pivot.minus is GND and cannot be part of the matrix (linear dependency)
        this.junctions.add(junction);
      }
      wires.addAll(junction.getWires());
    }

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

  private void fromSchematic() {
    wires.forEach(
      wire -> addResistor(wire.getA(), wire.getB(), NO_RESISTANCE)
    );
    parts.forEach(part -> {
      for (final InnerConnection innerConnection : part.getInnerConnections()) {
        addResistor(innerConnection.a, innerConnection.b, innerConnection.resistance);
      }
      if (part instanceof final PartBattery battery) {
        addBattery(battery);
      } else if (part instanceof final PartLight light) {
        addResistor(light.getA(), light.getB(), light.getResistance());
      }
    });
  }

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

  private double getPotential(final @NonNull Junction junction) {
    final int index = getJunctionIndex(junction);
    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 = getJunctionIndex(wire.getA());
    final int indexB = getJunctionIndex(wire.getB());
    final double conductance = indexA >= 0 && indexB >= 0 ? matrix.getEntry(indexA, indexB) : FULL_ADMITTANCE;
    final double potentialDifference = getPotential(wire.getB()) - getPotential(wire.getA());
    return conductance * potentialDifference;
  }

  private int getJunctionIndex(@NonNull final Junction junction) {
    for (int i = 0; i < junctions.size(); i++) {
      if (junctions.get(i) == junction) {
        return i;
      }
    }
    return -1;
  }

  public void addResistor(final Junction a, final Junction b, final double resistance) {
    final int indexA = getJunctionIndex(a);
    final int indexB = getJunctionIndex(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 addBattery(final PartBattery battery) {
    final double current = battery.getVoltage() / battery.getResistance();
    final int indexMinus = getJunctionIndex(battery.getMinus());
    final int indexPlus = getJunctionIndex(battery.getPlus());
    addCurrentSource(indexMinus, indexPlus, current, battery.getResistance());
  }

  public void addCurrentSource(final int index0, final int index1, final double current, final double resistance) {
    if (index0 >= 0) {
      currents.setEntry(index0, currents.getEntry(index0) - current);
    }
    if (index1 >= 0) {
      currents.setEntry(index1, currents.getEntry(index1) + current);
    }
    addResistor(index0, index1, resistance);
    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 int integer;

    private int decimal;

    private final String format;

    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);
    }

  }

}