/// GH: babylonian basins of attraction

#include <iostream>
#include <vector>
#include <cmath>
#include <fstream>
#include "Complx.h"
using namespace std;

const double pi = 3.1415926535897932384626433832795;

Complx cubeRoot(Complx, Complx);
Complx nthRoot(Complx, Complx, double);
Complx power(Complx, int);
void seeOut(vector<vector<unsigned> >, ofstream&);

int main() {
    double x1,x2,y1,y2,samples;
    cout << "\nThis program will map basins of attraction for the n roots of unity.";
    cout << "\nWhat value of n would you like? ";
    double n;
    cin >> n;
    cout << "\nWhat are the coordinates of the lower left corner of your window?\n";
    cin >> x1 >> y1;
    cout << "\nWhat are the coordinates of the upper right corner of your window?\n";
    cin >> x2 >> y2;
    cout << "\nWhat is the mesh size: ";
    cin >> samples;
    /// set up file for writing
    ofstream out("roots.txt");
    if(!out) cerr << "too bad, no file.";

    ///roots numbered 1,2,3 go here
    vector<vector<unsigned> > rootGrid;
    vector<unsigned> column; ///a member of rootGrid

    Complx initial{0,0}, converge, A{1,0};

    /// The three complex cube roots of i:
    const Complx root1{sqrt(3.)/2,0.5};
    const Complx root2{-sqrt(3.)/2,0.5};
    const Complx root3{0,-1};

    ///  The n nth roots of unity
    vector<Complx> roots;
    for(double i = 0; i < n; ++i) {
        roots.push_back(Complx(cos(2*pi*i/n),sin(2*pi*i/n)));
    }
    cout << "\nThe " << n << " roots of unity are:\n";
    for(double i = 0; i < n; ++i) {
        cout << roots[i] << endl;
    }

    ///for each column (starting at upper left corner)
    for(double x = x1; x <= x2; x += (x2-x1)/samples) {
        /// for each row
        column.clear();
        for(double y = y1; y <= y2; y += (y2-y1)/samples) {
            initial.re(x); initial.im(y);
            converge = nthRoot(A,initial,n);
            cout << converge;
            if(abs(converge-roots[0])<0.1) column.push_back(1);
            else if(abs(converge-roots[1])<0.1) column.push_back(2);
            else if(abs(converge-roots[2])<0.1) column.push_back(3);
            //else if(abs(converge-roots[3])<0.1) column.push_back(4);
            else column.push_back(5);
        }
        rootGrid.push_back(column);
    }
    seeOut(rootGrid, out);
}

Complx cubeRoot(Complx A, Complx temp) {
    ///Complx temp{A.real()/3,A.imag()/3};
    Complx next{0,0};
    next = (temp*2+A/(temp*temp))/3;
    while((next-temp).m()>A.m()*1e-10) {
        temp = next;
        next = (temp*2+A/(temp*temp))/3;
    }
    return next;
}

Complx power(Complx c, int pwr) {
    Complx result{1,0};
    while(pwr) {
        if(pwr & 1) { /// if pwr is odd
            result = c*result;
        }
        pwr >>= 1;  /// divide power in half
        c = c*c;  /// square c
    }
    return result;
}

Complx nthRoot(Complx A, Complx temp, double n) {
    ///Use babylonian algorithm to compute nth root of A
    ///Complx temp{A.real()/3,A.imag()/3};
    Complx next;
    next = (temp*(n-1)+A/power(temp,int(n-1)))/n;
    while(abs(next-temp)>abs(A)*1e-10) {
        temp = next;
        next = (temp*(n-1)+A/power(temp,int(n-1)))/n;
        ///cout << temp;
        ///cin.get();
    }
    return next;
}

void seeOut(vector<vector<unsigned> > v, ofstream& out) {
    //vector<unsigned> col;
    for(int i = 0; i < v.size(); ++i) {
        ///col = v[i];
        for(int j = 0; j < v[i].size(); ++j) {
            //cout << v[i][j];
            out << v[i][j] << " ";
        }
        ///cout << endl;
        out << endl;
    }
}